frp/test/e2e/framework/framework.go

267 lines
6.1 KiB
Go
Raw Normal View History

2020-06-02 22:48:55 +08:00
package framework
import (
"bytes"
"fmt"
"os"
2021-08-02 13:07:28 +08:00
"path/filepath"
2020-06-02 22:48:55 +08:00
"regexp"
"strings"
"text/template"
2023-02-27 14:44:16 +08:00
"github.com/onsi/ginkgo/v2"
2022-08-29 01:02:53 +08:00
"github.com/fatedier/frp/test/e2e/mock/server"
2020-06-02 22:48:55 +08:00
"github.com/fatedier/frp/test/e2e/pkg/port"
"github.com/fatedier/frp/test/e2e/pkg/process"
)
type Options struct {
TotalParallelNode int
CurrentNodeIndex int
FromPortIndex int
ToPortIndex int
}
type Framework struct {
TempDirectory string
2021-03-31 16:57:39 +08:00
// ports used in this framework indexed by port name.
usedPorts map[string]int
2023-11-22 14:30:22 +08:00
// record ports allocated by this framework and release them after each test
allocatedPorts []int
2021-06-18 16:48:36 +08:00
2021-03-31 16:57:39 +08:00
// portAllocator to alloc port for this test case.
2020-06-02 22:48:55 +08:00
portAllocator *port.Allocator
// Multiple default mock servers used for e2e testing.
2020-06-02 22:48:55 +08:00
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
// Manual registered mock servers.
2021-06-18 16:48:36 +08:00
servers []server.Server
// used to generate unique config file name.
configFileIndex int64
// envs used to start processes, the form is `key=value`.
osEnvs []string
2020-06-02 22:48:55 +08:00
}
func NewDefaultFramework() *Framework {
2023-02-27 14:44:16 +08:00
suiteConfig, _ := ginkgo.GinkgoConfiguration()
2020-06-02 22:48:55 +08:00
options := Options{
2023-02-27 14:44:16 +08:00
TotalParallelNode: suiteConfig.ParallelTotal,
CurrentNodeIndex: suiteConfig.ParallelProcess,
2023-05-29 00:27:27 +08:00
FromPortIndex: 10000,
2023-12-21 22:46:08 +08:00
ToPortIndex: 30000,
2020-06-02 22:48:55 +08:00
}
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),
2020-06-02 22:48:55 +08:00
}
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 := os.MkdirTemp(os.TempDir(), "frp-e2e-test-*")
2020-06-02 22:48:55 +08:00
ExpectNoError(err)
f.TempDirectory = dir
2020-09-07 14:57:23 +08:00
f.mockServers = NewMockServers(f.portAllocator)
if err := f.mockServers.Run(); err != nil {
Failf("%v", err)
}
2021-06-18 16:48:36 +08:00
params := f.mockServers.GetTemplateParams()
for k, v := range params {
switch t := v.(type) {
case int:
2022-08-29 01:02:53 +08:00
f.usedPorts[k] = t
default:
2021-06-18 16:48:36 +08:00
}
}
2020-06-02 22:48:55 +08:00
}
func (f *Framework) AfterEach() {
if !f.beforeEachStarted {
return
}
RemoveCleanupAction(f.cleanupHandle)
2020-09-07 14:57:23 +08:00
// stop processor
2020-06-02 22:48:55 +08:00
for _, p := range f.serverProcesses {
2022-08-29 01:02:53 +08:00
_ = p.Stop()
2023-05-29 00:27:27 +08:00
if TestContext.Debug || ginkgo.CurrentSpecReport().Failed() {
2020-09-07 14:57:23 +08:00
fmt.Println(p.ErrorOutput())
fmt.Println(p.StdOutput())
}
2020-06-02 22:48:55 +08:00
}
for _, p := range f.clientProcesses {
2022-08-29 01:02:53 +08:00
_ = p.Stop()
2023-05-29 00:27:27 +08:00
if TestContext.Debug || ginkgo.CurrentSpecReport().Failed() {
2020-09-07 14:57:23 +08:00
fmt.Println(p.ErrorOutput())
fmt.Println(p.StdOutput())
}
2020-06-02 22:48:55 +08:00
}
f.serverProcesses = nil
f.clientProcesses = nil
2020-09-07 14:57:23 +08:00
// close default mock servers
2020-09-07 14:57:23 +08:00
f.mockServers.Close()
// close manual registered mock servers
for _, s := range f.servers {
s.Close()
}
2020-09-07 14:57:23 +08:00
// clean directory
os.RemoveAll(f.TempDirectory)
f.TempDirectory = ""
2021-06-18 16:48:36 +08:00
f.serverConfPaths = []string{}
f.clientConfPaths = []string{}
2020-09-07 14:57:23 +08:00
// release used ports
2021-03-31 16:57:39 +08:00
for _, port := range f.usedPorts {
2020-09-07 14:57:23 +08:00
f.portAllocator.Release(port)
}
f.usedPorts = make(map[string]int)
2021-06-18 16:48:36 +08:00
2023-11-22 14:30:22 +08:00
// release allocated ports
for _, port := range f.allocatedPorts {
2021-06-18 16:48:36 +08:00
f.portAllocator.Release(port)
}
2023-11-22 14:30:22 +08:00
f.allocatedPorts = make([]int, 0)
2021-06-18 16:48:36 +08:00
// clear os envs
f.osEnvs = make([]string, 0)
2020-06-02 22:48:55 +08:00
}
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 {
2021-03-31 16:57:39 +08:00
port := f.portAllocator.GetByName(name)
2020-06-02 22:48:55 +08:00
if port <= 0 {
return nil, fmt.Errorf("can't allocate port")
}
ports[name] = port
}
return
}
2021-03-31 16:57:39 +08:00
// RenderTemplates alloc all ports for port names placeholder.
2020-06-02 22:48:55 +08:00
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
}
2020-06-02 22:48:55 +08:00
for _, t := range templates {
tmpl, err := template.New("frp-e2e").Parse(t)
2020-06-02 22:48:55 +08:00
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
}
2021-03-31 16:57:39 +08:00
func (f *Framework) PortByName(name string) int {
return f.usedPorts[name]
}
func (f *Framework) AllocPort() int {
port := f.portAllocator.Get()
ExpectTrue(port > 0, "alloc port failed")
2023-11-22 14:30:22 +08:00
f.allocatedPorts = append(f.allocatedPorts, port)
return port
}
func (f *Framework) ReleasePort(port int) {
f.portAllocator.Release(port)
}
2021-06-18 16:48:36 +08:00
func (f *Framework) RunServer(portName string, s server.Server) {
f.servers = append(f.servers, s)
2021-06-18 16:48:36 +08:00
if s.BindPort() > 0 && portName != "" {
f.usedPorts[portName] = s.BindPort()
}
err := s.Run()
2021-06-18 16:48:36 +08:00
ExpectNoError(err, "RunServer: with PortName %s", portName)
}
func (f *Framework) SetEnvs(envs []string) {
f.osEnvs = envs
}
2021-08-02 13:07:28 +08:00
func (f *Framework) WriteTempFile(name string, content string) string {
filePath := filepath.Join(f.TempDirectory, name)
2022-08-29 01:02:53 +08:00
err := os.WriteFile(filePath, []byte(content), 0o766)
2021-08-02 13:07:28 +08:00
ExpectNoError(err)
return filePath
}