package framework import ( "bytes" "fmt" "io/ioutil" "os" "regexp" "strings" "text/template" "github.com/fatedier/frp/test/e2e/pkg/port" "github.com/fatedier/frp/test/e2e/pkg/process" "github.com/onsi/ginkgo" "github.com/onsi/ginkgo/config" ) type Options struct { TotalParallelNode int CurrentNodeIndex int FromPortIndex int ToPortIndex int } type Framework struct { TempDirectory string UsedPorts map[string]int portAllocator *port.Allocator // Multiple mock servers used for e2e testing. mockServers *MockServers // To make sure that this framework cleans up after itself, no matter what, // we install a Cleanup action before each test and clear it after. If we // should abort, the AfterSuite hook should run all Cleanup actions. cleanupHandle CleanupActionHandle // beforeEachStarted indicates that BeforeEach has started beforeEachStarted bool serverConfPaths []string serverProcesses []*process.Process clientConfPaths []string clientProcesses []*process.Process } func NewDefaultFramework() *Framework { options := Options{ TotalParallelNode: config.GinkgoConfig.ParallelTotal, CurrentNodeIndex: config.GinkgoConfig.ParallelNode, FromPortIndex: 20000, ToPortIndex: 50000, } return NewFramework(options) } func NewFramework(opt Options) *Framework { f := &Framework{ portAllocator: port.NewAllocator(opt.FromPortIndex, opt.ToPortIndex, opt.TotalParallelNode, opt.CurrentNodeIndex-1), } ginkgo.BeforeEach(f.BeforeEach) ginkgo.AfterEach(f.AfterEach) return f } // BeforeEach create a temp directory. func (f *Framework) BeforeEach() { f.beforeEachStarted = true f.cleanupHandle = AddCleanupAction(f.AfterEach) dir, err := ioutil.TempDir(os.TempDir(), "frpe2e-test-*") ExpectNoError(err) f.TempDirectory = dir f.mockServers = NewMockServers(f.portAllocator) if err := f.mockServers.Run(); err != nil { Failf("%v", err) } } func (f *Framework) AfterEach() { if !f.beforeEachStarted { return } RemoveCleanupAction(f.cleanupHandle) // stop processor for _, p := range f.serverProcesses { p.Stop() if TestContext.Debug { fmt.Println(p.ErrorOutput()) fmt.Println(p.StdOutput()) } } for _, p := range f.clientProcesses { p.Stop() if TestContext.Debug { fmt.Println(p.ErrorOutput()) fmt.Println(p.StdOutput()) } } f.serverProcesses = nil f.clientProcesses = nil // close mock servers f.mockServers.Close() // clean directory os.RemoveAll(f.TempDirectory) f.TempDirectory = "" f.serverConfPaths = nil f.clientConfPaths = nil // release used ports for _, port := range f.UsedPorts { f.portAllocator.Release(port) } f.UsedPorts = nil } var portRegex = regexp.MustCompile(`{{ \.Port.*? }}`) // RenderPortsTemplate render templates with ports. // // Local: {{ .Port1 }} // Target: {{ .Port2 }} // // return rendered content and all allocated ports. 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.Get() if port <= 0 { return nil, fmt.Errorf("can't allocate port") } ports[name] = port } return } 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 _, t := range templates { tmpl, err := template.New("").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 }