ServerStatus/server/tcp_server_test.go

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"net"
	"strings"
	"testing"
	"time"
)

func TestAgentProtocolAndTrafficState(t *testing.T) {
	app := newTestApp(t, minimalTestConfig())
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	serveDone := make(chan error, 1)
	go func() { serveDone <- NewAgentServer(app).Serve(listener) }()

	connection, err := net.DialTimeout("tcp", listener.Addr().String(), time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer connection.Close()
	reader := bufio.NewReader(connection)
	readContains(t, reader, "Authentication required")
	if _, err := fmt.Fprintln(connection, "s01:secret"); err != nil {
		t.Fatal(err)
	}
	readContains(t, reader, "Authentication successful")
	readContains(t, reader, "You are connecting via: IPv4")
	readContains(t, reader, `"monitor":0`)

	if _, err := fmt.Fprintln(connection, "pong on"); err != nil {
		t.Fatal(err)
	}
	online6 := true
	update := AgentStats{
		Uptime: 90061, Load1: 1.25, Load5: 1, Load15: 0.75, CPU: 33.5, CPUCores: 4, CPUModel: "Test CPU",
		MemoryTotal: 1024, MemoryUsed: 512, SwapTotal: 128, SwapUsed: 2, HDDTotal: 10000, HDDUsed: 4000,
		NetworkRX: 123, NetworkTX: 456, NetworkIn: 1_000_000, NetworkOut: 2_000_000,
		Ping10010: 1, Ping189: 2, Ping10086: 3, Time10010: 10, Time189: 20, Time10086: 30,
		TCPCount: 10, UDPCount: 2, ProcessCount: 30, ThreadCount: 60, IORead: 1, IOWrite: 2,
		OS: "linux", Custom: "example=12", Online6: &online6,
	}
	payload, _ := json.Marshal(update)
	if _, err := fmt.Fprintf(connection, "update %s\n", payload); err != nil {
		t.Fatal(err)
	}
	readContains(t, reader, "0")

	eventually(t, time.Second, func() bool {
		servers := app.SnapshotStats()["servers"].([]any)
		server := servers[0].(map[string]any)
		return server["online4"] == true && server["online6"] == true && server["cpu_model"] == "Test CPU" && server["uptime"] == "1 天"
	})

	result, apiErr := app.ResetTraffic("s01")
	if apiErr != nil {
		t.Fatal(apiErr)
	}
	stats := result["stats"].(map[string]any)
	if stats["last_network_in"] != int64(1_000_000) {
		t.Fatalf("traffic reset result: %#v", stats)
	}

	duplicate, err := net.DialTimeout("tcp", listener.Addr().String(), time.Second)
	if err != nil {
		t.Fatal(err)
	}
	duplicateReader := bufio.NewReader(duplicate)
	readContains(t, duplicateReader, "Authentication required")
	_, _ = fmt.Fprintln(duplicate, "s01:secret")
	readContains(t, duplicateReader, "Only one connection per user")
	_ = duplicate.Close()

	wrong, err := net.DialTimeout("tcp", listener.Addr().String(), time.Second)
	if err != nil {
		t.Fatal(err)
	}
	wrongReader := bufio.NewReader(wrong)
	readContains(t, wrongReader, "Authentication required")
	_, _ = fmt.Fprintln(wrong, "s01:wrong")
	readContains(t, wrongReader, "Wrong username")
	_ = wrong.Close()

	if apiErr := app.ReloadConfig(); apiErr != nil {
		t.Fatal(apiErr)
	}
	readContains(t, reader, "Server reloading")
	eventually(t, time.Second, func() bool {
		server := app.SnapshotStats()["servers"].([]any)[0].(map[string]any)
		return server["online4"] == false && server["online6"] == false
	})

	app.cancel()
	_ = listener.Close()
	select {
	case err := <-serveDone:
		if err != nil {
			t.Fatal(err)
		}
	case <-time.After(time.Second):
		t.Fatal("agent server did not stop")
	}
}

func readContains(t *testing.T, reader *bufio.Reader, expected string) string {
	t.Helper()
	line, err := reader.ReadString('\n')
	if err != nil {
		t.Fatalf("read %q: %v", expected, err)
	}
	if !strings.Contains(line, expected) {
		t.Fatalf("expected %q in %q", expected, line)
	}
	return line
}