90210

test/e2e/framework/framework.go

@@ -24,58 +24,38 @@ type Options struct {
 }
 
 type Framework struct {
-	TempDirectory string
-
-	// פורטים בשימוש במסגרת הזו לפי שם.
-	usedPorts map[string]int
-
-	// רישום הפורטים שהוקצו וישוחררו אחרי כל בדיקה.
-	allocatedPorts []int
-
-	// הקצאת פורטים למקרה בדיקה זה.
-	portAllocator *port.Allocator
-
-	// מספר שרתי הדמיה (mock) עבור בדיקות קצה.
-	mockServers *MockServers
-
-	// כדי לוודא שהמסגרת הזו מנקה אחריה, לא משנה מה קורה,
-	// אנחנו מתקינים פעולה לניקוי לפני כל בדיקה ומסירים אותה אחרי.
-	// במקרה של כישלון, הפונקציה AfterSuite תפעיל את כל פעולות הניקוי.
-	cleanupHandle CleanupActionHandle
-
-	// מציין ש-BeforeEach התחיל
+	TempDirectory     string
+	usedPorts         map[string]int
+	allocatedPorts    []int
+	portAllocator     *port.Allocator
+	mockServers       *MockServers
+	cleanupHandle     CleanupActionHandle
 	beforeEachStarted bool
 
-	serverConfPaths []string
-	serverProcesses []*process.Process
-	clientConfPaths []string
-	clientProcesses []*process.Process
+	serverConfPaths   []string
+	serverProcesses   []*process.Process
+	clientConfPaths   []string
+	clientProcesses   []*process.Process
 
-	// שרתי mock שנרשמו באופן ידני.
-	servers []server.Server
-
-	// משמש ליצירת שם ייחודי לקובץ קונפיגורציה.
+	servers         []server.Server
 	configFileIndex int64
-
-	// משתני סביבה לצורך הפעלת תהליכים, בפורמט `key=value`.
-	osEnvs []string
+	osEnvs          []string
 }
 
 func NewDefaultFramework() *Framework {
-	suiteConfig, _ := ginkgo.GinkgoConfiguration()
-	options := Options{
-		TotalParallelNode: suiteConfig.ParallelTotal,
-		CurrentNodeIndex:  suiteConfig.ParallelProcess,
+	suiteCfg, _ := ginkgo.GinkgoConfiguration()
+	return NewFramework(Options{
+		TotalParallelNode: suiteCfg.ParallelTotal,
+		CurrentNodeIndex:  suiteCfg.ParallelProcess,
 		FromPortIndex:     10000,
 		ToPortIndex:       30000,
-	}
-	return NewFramework(options)
+	})
 }
 
 func NewFramework(opt Options) *Framework {
 	f := &Framework{
-		portAllocator: port.NewAllocator(opt.FromPortIndex, opt.ToPortIndex, opt.TotalParallelNode, opt.CurrentNodeIndex-1),
 		usedPorts:     make(map[string]int),
+		portAllocator: port.NewAllocator(opt.FromPortIndex, opt.ToPortIndex, opt.TotalParallelNode, opt.CurrentNodeIndex-1),
 	}
 
 	ginkgo.BeforeEach(f.BeforeEach)
@@ -83,10 +63,9 @@ func NewFramework(opt Options) *Framework {
 	return f
 }
 
-// BeforeEach יוצר תיקייה זמנית.
+// BeforeEach sets up temp directory, mock servers, and allocates base ports.
 func (f *Framework) BeforeEach() {
 	f.beforeEachStarted = true
-
 	f.cleanupHandle = AddCleanupAction(f.AfterEach)
 
 	dir, err := os.MkdirTemp(os.TempDir(), "frp-e2e-test-*")
@@ -94,149 +73,83 @@ func (f *Framework) BeforeEach() {
 	f.TempDirectory = dir
 
 	f.mockServers = NewMockServers(f.portAllocator)
-	if err := f.mockServers.Run(); err != nil {
-		Failf("%v", err)
-	}
+	ExpectNoError(f.mockServers.Run())
 
-	params := f.mockServers.GetTemplateParams()
-	for k, v := range params {
-		switch t := v.(type) {
-		case int:
-			f.usedPorts[k] = t
-		default:
+	for name, val := range f.mockServers.GetTemplateParams() {
+		if port, ok := val.(int); ok {
+			f.usedPorts[name] = port
 		}
 	}
 }
 
+// AfterEach performs cleanup: stopping processes, removing temp files, releasing ports.
 func (f *Framework) AfterEach() {
 	if !f.beforeEachStarted {
 		return
 	}
 
 	RemoveCleanupAction(f.cleanupHandle)
+	f.stopProcesses()
+	f.cleanupTempDirectory()
+	f.releaseAllPorts()
+	f.osEnvs = nil
+}
 
-	// עצירת התהליכים
-	for _, p := range f.serverProcesses {
-		_ = p.Stop()
-		if TestContext.Debug || ginkgo.CurrentSpecReport().Failed() {
-			fmt.Println(p.ErrorOutput())
-			fmt.Println(p.StdOutput())
-		}
+func (f *Framework) stopProcesses() {
+	for _, proc := range f.serverProcesses {
+		_ = proc.Stop()
+		printDebugOutput(proc)
 	}
-	for _, p := range f.clientProcesses {
-		_ = p.Stop()
-		if TestContext.Debug || ginkgo.CurrentSpecReport().Failed() {
-			fmt.Println(p.ErrorOutput())
-			fmt.Println(p.StdOutput())
-		}
+	for _, proc := range f.clientProcesses {
+		_ = proc.Stop()
+		printDebugOutput(proc)
 	}
 	f.serverProcesses = nil
 	f.clientProcesses = nil
 
-	// סגירת שרתי mock ברירת מחדל
 	f.mockServers.Close()
-
-	// סגירת שרתי mock שנרשמו ידנית
 	for _, s := range f.servers {
 		s.Close()
 	}
+}
 
-	// ניקוי התיקייה הזמנית
-	os.RemoveAll(f.TempDirectory)
-	f.TempDirectory = ""
-	f.serverConfPaths = []string{}
-	f.clientConfPaths = []string{}
+func printDebugOutput(p *process.Process) {
+	if TestContext.Debug || ginkgo.CurrentSpecReport().Failed() {
+		fmt.Println(p.ErrorOutput())
+		fmt.Println(p.StdOutput())
+	}
+}
 
-	// שחרור פורטים בשימוש
+func (f *Framework) cleanupTempDirectory() {
+	if f.TempDirectory != "" {
+		_ = os.RemoveAll(f.TempDirectory)
+		f.TempDirectory = ""
+	}
+	f.serverConfPaths = nil
+	f.clientConfPaths = nil
+}
+
+func (f *Framework) releaseAllPorts() {
 	for _, port := range f.usedPorts {
 		f.portAllocator.Release(port)
 	}
 	f.usedPorts = make(map[string]int)
 
-	// שחרור פורטים שהוקצו
 	for _, port := range f.allocatedPorts {
 		f.portAllocator.Release(port)
 	}
-	f.allocatedPorts = make([]int, 0)
-
-	// ניקוי משתני סביבה
-	f.osEnvs = make([]string, 0)
-}
-
-var portRegex = regexp.MustCompile(`{{ \.Port.*? }}`)
-
-// RenderPortsTemplate מרנדר תבניות עם פורטים.
-//
-// מקומי: {{ .Port1 }}
-// יעד:   {{ .Port2 }}
-//
-// מחזיר תוכן מרונדר וכל הפורטים שהוקצו.
-func (f *Framework) genPortsFromTemplates(templates []string) (ports map[string]int, err error) {
-	ports = make(map[string]int)
-	for _, t := range templates {
-		arrs := portRegex.FindAllString(t, -1)
-		for _, str := range arrs {
-			str = strings.TrimPrefix(str, "{{ .")
-			str = strings.TrimSuffix(str, " }}")
-			str = strings.TrimSpace(str)
-			ports[str] = 0
-		}
-	}
-	defer func() {
-		if err != nil {
-			for _, port := range ports {
-				f.portAllocator.Release(port)
-			}
-		}
-	}()
-
-	for name := range ports {
-		port := f.portAllocator.GetByName(name)
-		if port <= 0 {
-			return nil, fmt.Errorf("לא ניתן להקצות פורט")
-		}
-		ports[name] = port
-	}
-	return
-}
-
-// RenderTemplates מקצה פורטים לכל תבנית עם שמות פורטים.
-func (f *Framework) RenderTemplates(templates []string) (outs []string, ports map[string]int, err error) {
-	ports, err = f.genPortsFromTemplates(templates)
-	if err != nil {
-		return
-	}
-
-	params := f.mockServers.GetTemplateParams()
-	for name, port := range ports {
-		params[name] = port
-	}
-
-	for name, port := range f.usedPorts {
-		params[name] = port
-	}
-
-	for _, t := range templates {
-		tmpl, err := template.New("frp-e2e").Parse(t)
-		if err != nil {
-			return nil, nil, err
-		}
-		buffer := bytes.NewBuffer(nil)
-		if err = tmpl.Execute(buffer, params); err != nil {
-			return nil, nil, err
-		}
-		outs = append(outs, buffer.String())
-	}
-	return
+	f.allocatedPorts = nil
 }
 
+// PortByName returns the allocated port for a given name.
 func (f *Framework) PortByName(name string) int {
 	return f.usedPorts[name]
 }
 
+// AllocPort reserves a new port and tracks it for cleanup.
 func (f *Framework) AllocPort() int {
 	port := f.portAllocator.Get()
-	ExpectTrue(port > 0, "הקצאת פורט נכשלה")
+	ExpectTrue(port > 0, "failed to allocate port")
 	f.allocatedPorts = append(f.allocatedPorts, port)
 	return port
 }
@@ -245,22 +158,82 @@ func (f *Framework) ReleasePort(port int) {
 	f.portAllocator.Release(port)
 }
 
+// RunServer registers and starts a mock server.
 func (f *Framework) RunServer(portName string, s server.Server) {
 	f.servers = append(f.servers, s)
-	if s.BindPort() > 0 && portName != "" {
-		f.usedPorts[portName] = s.BindPort()
+	if port := s.BindPort(); port > 0 && portName != "" {
+		f.usedPorts[portName] = port
 	}
-	err := s.Run()
-	ExpectNoError(err, "RunServer: עם שם פורט %s", portName)
+	ExpectNoError(s.Run(), "RunServer failed for portName %s", portName)
 }
 
 func (f *Framework) SetEnvs(envs []string) {
 	f.osEnvs = envs
 }
 
-func (f *Framework) WriteTempFile(name string, content string) string {
-	filePath := filepath.Join(f.TempDirectory, name)
-	err := os.WriteFile(filePath, []byte(content), 0o600)
-	ExpectNoError(err)
-	return filePath
+// WriteTempFile writes `content` into a temp file under the framework's temp dir.
+func (f *Framework) WriteTempFile(name, content string) string {
+	path := filepath.Join(f.TempDirectory, name)
+	ExpectNoError(os.WriteFile(path, []byte(content), 0o600))
+	return path
+}
+
+// RenderTemplates renders Go templates and allocates required ports.
+func (f *Framework) RenderTemplates(templates []string) ([]string, map[string]int, error) {
+	ports, err := f.extractAndAllocPorts(templates)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	params := f.mockServers.GetTemplateParams()
+	for k, v := range ports {
+		params[k] = v
+	}
+	for k, v := range f.usedPorts {
+		params[k] = v
+	}
+
+	var rendered []string
+	for _, tmpl := range templates {
+		t, err := template.New("frp-e2e").Parse(tmpl)
+		if err != nil {
+			return nil, nil, err
+		}
+		var buf bytes.Buffer
+		if err := t.Execute(&buf, params); err != nil {
+			return nil, nil, err
+		}
+		rendered = append(rendered, buf.String())
+	}
+	return rendered, ports, nil
+}
+
+var portRegex = regexp.MustCompile(`{{\s*\.Port[^}]*}}`)
+
+// extractAndAllocPorts parses templates for named ports and allocates them.
+func (f *Framework) extractAndAllocPorts(templates []string) (map[string]int, error) {
+	ports := make(map[string]int)
+	for _, t := range templates {
+		for _, match := range portRegex.FindAllString(t, -1) {
+			name := strings.TrimSpace(strings.TrimSuffix(strings.TrimPrefix(match, "{{ ."), "}}"))
+			ports[name] = 0
+		}
+	}
+
+	defer func() {
+		if r := recover(); r != nil {
+			for _, port := range ports {
+				f.portAllocator.Release(port)
+			}
+		}
+	}()
+
+	for name := range ports {
+		p := f.portAllocator.GetByName(name)
+		if p <= 0 {
+			return nil, fmt.Errorf("failed to allocate port for %q", name)
+		}
+		ports[name] = p
+	}
+	return ports, nil
 }