new e2e framework (#1835)

test/e2e/framework/framework.go

@@ -0,0 +1,167 @@
+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),
+	}
+	f.mockServers = NewMockServers(f.portAllocator)
+	if err := f.mockServers.Run(); err != nil {
+		Failf("%v", err)
+	}
+
+	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
+}
+
+func (f *Framework) AfterEach() {
+	if !f.beforeEachStarted {
+		return
+	}
+
+	RemoveCleanupAction(f.cleanupHandle)
+
+	os.RemoveAll(f.TempDirectory)
+	f.TempDirectory = ""
+	f.UsedPorts = nil
+	f.serverConfPaths = nil
+	f.clientConfPaths = nil
+	for _, p := range f.serverProcesses {
+		p.Stop()
+	}
+	for _, p := range f.clientProcesses {
+		p.Stop()
+	}
+	f.serverProcesses = nil
+	f.clientProcesses = 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
+}