proto-socket/go/crosstest/go_python.go

//go:build ignore

package main

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"sync"
	"time"

	"google.golang.org/protobuf/proto"
	"nhooyr.io/websocket"

	toki "toki-labs.com/toki_socket/go"
	"toki-labs.com/toki_socket/go/packets"
)

const (
	host                    = "127.0.0.1"
	goPythonTCPPort         = 29390
	goPythonWSPort          = 29392
	wsPath                  = "/"
	processTimeout          = 20 * time.Second
	serverObservationWindow = 200 * time.Millisecond
)

func parserMap() toki.ParserMap {
	return toki.ParserMap{
		toki.TypeNameOf(&packets.TestData{}): func(b []byte) (proto.Message, error) {
			m := &packets.TestData{}
			return m, proto.Unmarshal(b, m)
		},
	}
}

func main() {
	fmt.Printf("INFO typeName go=%s\n", toki.TypeNameOf(&packets.TestData{}))
	if err := run(); err != nil {
		fmt.Fprintf(os.Stderr, "FAIL crosstest error=%v\n", err)
		os.Exit(1)
	}
	fmt.Println("PASS all go-server/python-client crosstests passed")
}

func run() error {
	if err := runTCPSendPush(); err != nil {
		return err
	}
	time.Sleep(150 * time.Millisecond)
	if err := runTCPRequests(); err != nil {
		return err
	}
	time.Sleep(150 * time.Millisecond)
	if err := runWSSendPush(); err != nil {
		return err
	}
	time.Sleep(150 * time.Millisecond)
	return runWSRequests()
}

func runTCPSendPush() error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	received := make(chan bool, 1)
	server := toki.NewTcpServer(host, goPythonTCPPort, func(conn net.Conn) *toki.TcpClient {
		return toki.NewTcpClient(conn, 0, 0, parserMap())
	})
	server.OnClientConnected = func(client *toki.TcpClient) {
		toki.AddListenerTyped[*packets.TestData](&client.Communicator, func(data *packets.TestData) {
			fmt.Printf("SERVER_RECEIVED index=%d message=%s\n", data.GetIndex(), data.GetMessage())
			valid := data.GetIndex() == 101 && data.GetMessage() == "fire from python client"
			select {
			case received <- valid:
			default:
			}
			if valid {
				_ = client.Send(&packets.TestData{Index: 200, Message: "push from go server"})
			}
		})
	}
	if err := server.Start(ctx); err != nil {
		return err
	}
	defer server.Stop()

	if err := runPythonClient("tcp", goPythonTCPPort, "send-push", map[string]bool{"1": true, "2": true}); err != nil {
		return err
	}
	select {
	case ok := <-received:
		if !ok {
			return errors.New("TCP send-push server received unexpected data")
		}
	case <-time.After(serverObservationWindow):
		return errors.New("TCP send-push server did not receive expected data")
	}
	return nil
}

func runTCPRequests() error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	server := toki.NewTcpServer(host, goPythonTCPPort, func(conn net.Conn) *toki.TcpClient {
		return toki.NewTcpClient(conn, 0, 0, parserMap())
	})
	server.OnClientConnected = func(client *toki.TcpClient) {
		toki.AddRequestListenerTyped[*packets.TestData, *packets.TestData](&client.Communicator, func(req *packets.TestData) (*packets.TestData, error) {
			return &packets.TestData{
				Index:   req.GetIndex() * 2,
				Message: "echo: " + req.GetMessage(),
			}, nil
		})
	}
	if err := server.Start(ctx); err != nil {
		return err
	}
	defer server.Stop()

	return runPythonClient("tcp", goPythonTCPPort, "requests", map[string]bool{"3": true, "4": true})
}

func runWSSendPush() error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	received := make(chan bool, 1)
	server := toki.NewWsServer(host, goPythonWSPort, wsPath, func(conn *websocket.Conn) *toki.WsClient {
		return toki.NewWsClient(conn, 0, 0, parserMap())
	})
	server.OnClientConnected = func(client *toki.WsClient) {
		toki.AddListenerTyped[*packets.TestData](&client.Communicator, func(data *packets.TestData) {
			fmt.Printf("SERVER_RECEIVED index=%d message=%s\n", data.GetIndex(), data.GetMessage())
			valid := data.GetIndex() == 101 && data.GetMessage() == "fire from python client"
			select {
			case received <- valid:
			default:
			}
			if valid {
				_ = client.Send(&packets.TestData{Index: 200, Message: "push from go server"})
			}
		})
	}
	if err := server.Start(ctx); err != nil {
		return err
	}
	defer server.Stop()

	if err := runPythonClient("ws", goPythonWSPort, "send-push", map[string]bool{"1": true, "2": true}); err != nil {
		return err
	}
	select {
	case ok := <-received:
		if !ok {
			return errors.New("WS send-push server received unexpected data")
		}
	case <-time.After(serverObservationWindow):
		return errors.New("WS send-push server did not receive expected data")
	}
	return nil
}

func runWSRequests() error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	server := toki.NewWsServer(host, goPythonWSPort, wsPath, func(conn *websocket.Conn) *toki.WsClient {
		return toki.NewWsClient(conn, 0, 0, parserMap())
	})
	server.OnClientConnected = func(client *toki.WsClient) {
		toki.AddRequestListenerTyped[*packets.TestData, *packets.TestData](&client.Communicator, func(req *packets.TestData) (*packets.TestData, error) {
			return &packets.TestData{
				Index:   req.GetIndex() * 2,
				Message: "echo: " + req.GetMessage(),
			}, nil
		})
	}
	if err := server.Start(ctx); err != nil {
		return err
	}
	defer server.Stop()

	return runPythonClient("ws", goPythonWSPort, "requests", map[string]bool{"3": true, "4": true})
}

func runPythonClient(mode string, port int, phase string, expected map[string]bool) error {
	ctx, cancel := context.WithTimeout(context.Background(), processTimeout)
	defer cancel()

	pythonDir, err := pythonPackageDir()
	if err != nil {
		return err
	}

	cmd := exec.CommandContext(ctx, "python3", "crosstest/go_python_client.py",
		"--mode="+mode,
		fmt.Sprintf("--port=%d", port),
		"--phase="+phase,
	)
	cmd.Dir = pythonDir

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return err
	}
	if err := cmd.Start(); err != nil {
		return err
	}

	var mu sync.Mutex
	var resultLines []string
	var wg sync.WaitGroup
	wg.Add(2)
	go scanLines(&wg, stdoutPipe, os.Stdout, &mu, &resultLines)
	go scanLines(&wg, stderrPipe, os.Stderr, nil, nil)

	waitErr := cmd.Wait()
	wg.Wait()
	if ctx.Err() == context.DeadlineExceeded {
		return fmt.Errorf("python client %s/%s timed out", mode, phase)
	}
	return validateResultLines("python-client "+mode+"/"+phase, waitErr, resultLines, expected)
}

func pythonPackageDir() (string, error) {
	candidates := make([]string, 0, 3)

	_, filename, _, ok := runtime.Caller(0)
	if ok {
		repoRoot := filepath.Dir(filepath.Dir(filepath.Dir(filename)))
		candidates = append(candidates, filepath.Join(repoRoot, "python"))
	}
	if wd, err := os.Getwd(); err == nil {
		candidates = append(candidates, findPythonPackageCandidates(wd)...)
	}
	if executable, err := os.Executable(); err == nil {
		candidates = append(candidates, findPythonPackageCandidates(filepath.Dir(executable))...)
	}

	for _, candidate := range candidates {
		if isPythonPackageDir(candidate) {
			return candidate, nil
		}
	}
	return "", fmt.Errorf("cannot resolve python package directory from candidates %v", candidates)
}

func findPythonPackageCandidates(start string) []string {
	candidates := make([]string, 0)
	for dir := start; ; dir = filepath.Dir(dir) {
		candidates = append(candidates, filepath.Join(dir, "python"))
		if filepath.Base(dir) == "python" {
			candidates = append(candidates, dir)
		}
		parent := filepath.Dir(dir)
		if parent == dir {
			return candidates
		}
	}
}

func isPythonPackageDir(dir string) bool {
	info, err := os.Stat(filepath.Join(dir, "pyproject.toml"))
	return err == nil && !info.IsDir()
}

func scanLines(wg *sync.WaitGroup, reader io.Reader, writer *os.File, mu *sync.Mutex, resultLines *[]string) {
	defer wg.Done()
	scanner := bufio.NewScanner(reader)
	for scanner.Scan() {
		line := scanner.Text()
		fmt.Fprintln(writer, line)
		if mu != nil && startsWithResult(line) {
			mu.Lock()
			*resultLines = append(*resultLines, line)
			mu.Unlock()
		}
	}
}

func validateResultLines(label string, waitErr error, lines []string, expected map[string]bool) error {
	failed := make([]string, 0)
	passed := make(map[string]bool)
	re := regexp.MustCompile(`scenario=([^ ]+)`)
	for _, line := range lines {
		if len(line) >= 5 && line[:5] == "FAIL " {
			failed = append(failed, line)
			continue
		}
		if len(line) >= 5 && line[:5] == "PASS " {
			match := re.FindStringSubmatch(line)
			if len(match) == 2 {
				passed[match[1]] = true
			}
		}
	}

	missing := make([]string, 0)
	for scenario := range expected {
		if !passed[scenario] {
			missing = append(missing, scenario)
		}
	}
	if waitErr != nil || len(failed) > 0 || len(missing) > 0 {
		return fmt.Errorf("%s failed waitErr=%v failed=%v missing=%v", label, waitErr, failed, missing)
	}
	return nil
}

func startsWithResult(line string) bool {
	return (len(line) >= 5 && line[:5] == "PASS ") || (len(line) >= 5 && line[:5] == "FAIL ")
}