frp/test/e2e/framework/mockservers.go

package framework

import (
	"github.com/fatedier/frp/test/e2e/mock/echoserver"
	"github.com/fatedier/frp/test/e2e/pkg/port"
)

const (
	TCPEchoServerPort = "TCPEchoServerPort"
	UDPEchoServerPort = "UDPEchoServerPort"
)

type MockServers struct {
	tcpEchoServer *echoserver.Server
	udpEchoServer *echoserver.Server
}

func NewMockServers(portAllocator *port.Allocator) *MockServers {
	s := &MockServers{}
	tcpPort := portAllocator.Get()
	udpPort := portAllocator.Get()
	s.tcpEchoServer = echoserver.New(echoserver.Options{
		Type:      echoserver.TCP,
		BindAddr:  "127.0.0.1",
		BindPort:  int32(tcpPort),
		RepeatNum: 1,
	})
	s.udpEchoServer = echoserver.New(echoserver.Options{
		Type:      echoserver.UDP,
		BindAddr:  "127.0.0.1",
		BindPort:  int32(udpPort),
		RepeatNum: 1,
	})
	return s
}

func (m *MockServers) Run() error {
	if err := m.tcpEchoServer.Run(); err != nil {
		return err
	}
	if err := m.udpEchoServer.Run(); err != nil {
		return err
	}
	return nil
}

func (m *MockServers) GetTemplateParams() map[string]interface{} {
	ret := make(map[string]interface{})
	ret[TCPEchoServerPort] = m.tcpEchoServer.GetOptions().BindPort
	ret[UDPEchoServerPort] = m.udpEchoServer.GetOptions().BindPort
	return ret
}

func (m *MockServers) GetParam(key string) interface{} {
	params := m.GetTemplateParams()
	if v, ok := params[key]; ok {
		return v
	}
	return nil
}