creator-hub/internal/controlplane/api/gateway_channel_test.go

package api

import (
	"context"
	"encoding/json"
	"errors"
	"testing"
	"time"

	hub "git.ipao.vip/rogee/creator-hub/internal/environment"
	"github.com/fasthttp/websocket"
)

func TestGatewayChannelDispatchesTaskAndWaitsForResult(t *testing.T) {
	store := newMemoryChannelStore(newTestGateway(1, "gw-main", "test-access-key-0001"))
	server := (&gatewayChannelTestServer{store: store, URL: startChannelServer(t, store)})

	SetGatewayChannelStore(store)

	fake := connectFakeGateway(t, server.URL, "test-access-key-0001", func(task taskFrame) *TaskResult {
		if task.Method != "POST" || task.Path != "/v1/browsers/alias/douyin/get" {
			return &TaskResult{Status: 404, Error: "unexpected path " + task.Path}
		}
		return &TaskResult{Status: 200, Body: []byte(`{"ok":true}`)}
	})
	_ = fake
	<-waitForSession(t, 1)

	status, body, err := gatewayCall(context.Background(), newTestGateway(1, "gw-main", "unused"), "POST", "/v1/browsers/alias/douyin/get", map[string]any{"url": "https://x"}, 5*time.Second)
	if err != nil || status != 200 || string(body) != `{"ok":true}` {
		t.Fatalf("gateway call failed: status=%d body=%s err=%v", status, body, err)
	}
}

func waitForSession(t *testing.T, gatewayID int64) <-chan struct{} {
	t.Helper()
	done := make(chan struct{})
	go func() {
		defer close(done)
		deadline := time.Now().Add(2 * time.Second)
		for time.Now().Before(deadline) {
			if sessionByGateway(gatewayID) != nil {
				return
			}
			time.Sleep(5 * time.Millisecond)
		}
		t.Error("gateway session never registered")
	}()
	return done
}

func TestGatewayChannelRejectsInvalidAccessKey(t *testing.T) {
	store := newMemoryChannelStore(newTestGateway(1, "gw-main", "test-access-key-0001"))
	server := (&gatewayChannelTestServer{store: store, URL: startChannelServer(t, store)})

	header := map[string][]string{"Authorization": {"Bearer wrong-key"}}
	conn, _, err := websocket.DefaultDialer.Dial(serverWSURL(server.URL)+"/v1/agent", header)
	if err != nil {
		t.Fatalf("dial before auth rejection is transport level: %v", err)
	}
	defer conn.Close()
	// 服务端应立即关闭连接：读操作收尾即失败。
	_ = conn.SetReadDeadline(time.Now().Add(2 * time.Second))
	if _, _, err = conn.ReadMessage(); err == nil {
		t.Fatal("invalid access key must not establish a session")
	}
	if sessionByGateway(1) != nil {
		t.Fatal("rejected gateway must not register a session")
	}
}

func TestGatewayChannelOfflineFailsSyncTaskImmediately(t *testing.T) {
	store := newMemoryChannelStore(newTestGateway(2, "gw-offline", "test-access-key-0002"))
	SetGatewayChannelStore(store)

	_, _, err := gatewayCall(context.Background(), newTestGateway(2, "gw-offline", "unused"), "POST", "/v1/browsers", nil, time.Second)
	if !errors.Is(err, hub.ErrGatewayOffline) {
		t.Fatalf("offline gateway must fail immediately: %v", err)
	}
	tasks, _ := store.NextPendingGatewayTasks(context.Background(), 2, 10)
	if len(tasks) != 0 {
		t.Fatalf("offline sync task must be terminal, pending=%d", len(tasks))
	}
}

func TestGatewayChannelTaskTimeoutMarksFailure(t *testing.T) {
	store := newMemoryChannelStore(newTestGateway(3, "gw-slow", "test-access-key-0003"))
	server := (&gatewayChannelTestServer{store: store, URL: startChannelServer(t, store)})
	SetGatewayChannelStore(store)

	// 假网关收到任务后不回结果。
	_ = connectFakeGateway(t, server.URL, "test-access-key-0003", func(task taskFrame) *TaskResult { return nil })
	<-waitForSession(t, 3)

	_, _, err := gatewayCall(context.Background(), newTestGateway(3, "gw-slow", "unused"), "POST", "/v1/browsers/alias/stop", nil, 300*time.Millisecond)
	if err == nil {
		t.Fatal("unanswered task must time out")
	}
	var offline hub.GatewayTask
	found := false
	for _, task := range func() map[int64]*hub.GatewayTask {
		store.mu.Lock()
		defer store.mu.Unlock()
		return store.tasks
	}() {
		if task.Error == "网关任务等待超时" {
			offline = *task
			found = true
		}
	}
	if !found || offline.Status != hub.TaskFailed {
		t.Fatalf("timeout must be recorded as failed task: %+v", offline)
	}
}

func TestGatewayChannelDisconnectInterruptsInflightTask(t *testing.T) {
	store := newMemoryChannelStore(newTestGateway(4, "gw-drop", "test-access-key-0004"))
	server := (&gatewayChannelTestServer{store: store, URL: startChannelServer(t, store)})
	SetGatewayChannelStore(store)

	fake := connectFakeGateway(t, server.URL, "test-access-key-0004", func(task taskFrame) *TaskResult { return nil })
	<-waitForSession(t, 4)

	callErr := make(chan error, 1)
	go func() {
		_, _, err := gatewayCall(context.Background(), newTestGateway(4, "gw-drop", "unused"), "POST", "/v1/browsers/alias/start", nil, 10*time.Second)
		callErr <- err
	}()
	deadline := time.After(2 * time.Second)
	var taskID int64
	select {
	case taskID = <-fake.taskIDs:
	case <-deadline:
		t.Fatal("task was never dispatched")
	}
	// 网关断开：等待方必须以连接中断收尾。
	_ = fake.conn.Close()
	select {
	case err := <-callErr:
		if err == nil {
			t.Fatal("disconnect must fail the waiting task")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("waiting task did not finish after disconnect")
	}
	_ = taskID
	deadline = time.After(2 * time.Second)
	for {
		store.mu.Lock()
		interrupted := store.interrupted
		store.mu.Unlock()
		if interrupted >= 1 {
			break
		}
		select {
		case <-deadline:
			t.Fatal("disconnect must interrupt dispatched tasks")
		default:
			time.Sleep(5 * time.Millisecond)
		}
	}
}

func TestGatewayChannelEventRoutingAndAck(t *testing.T) {
	store := newMemoryChannelStore(newTestGateway(5, "gw-events", "test-access-key-0005"))
	server := (&gatewayChannelTestServer{store: store, URL: startChannelServer(t, store)})

	acks := make(chan []string, 4)
	var conn *websocket.Conn
	ackReady := make(chan struct{})
	go func() {
		header := map[string][]string{"Authorization": {"Bearer test-access-key-0005"}}
		c, _, err := websocket.DefaultDialer.Dial(serverWSURL(server.URL)+"/v1/agent", header)
		if err != nil {
			t.Errorf("fake gateway dial failed: %v", err)
			return
		}
		conn = c
		go func() {
			for {
				var frame struct {
					Type        string   `json:"type"`
					Alias       string   `json:"alias"`
					DeliveryIDs []string `json:"delivery_ids"`
				}
				if err := c.ReadJSON(&frame); err != nil {
					return
				}
				if frame.Type == "event_ack" {
					acks <- frame.DeliveryIDs
				}
			}
		}()
		close(ackReady)
	}()
	<-ackReady
	defer func() { _ = conn.Close() }()
	<-waitForSession(t, 5)
	session := sessionByGateway(5)
	queue := session.subscribeEvents("alias-a")
	defer session.unsubscribeEvents("alias-a")

	// 平台侧尚未消费前，网关先推两批：队列应缓冲而非丢弃。
	batch1 := eventBatchFor("alias-a", "d-1")
	batch2 := eventBatchFor("alias-a", "d-2")
	if err := conn.WriteJSON(map[string]any{"type": "event", "event": batch1}); err != nil {
		t.Fatal(err)
	}
	if err := conn.WriteJSON(map[string]any{"type": "event", "event": batch2}); err != nil {
		t.Fatal(err)
	}

	got := <-queue
	if got.Deliveries[0].DeliveryID != "d-1" {
		t.Fatalf("first batch lost: %+v", got)
	}
	if err := session.acknowledgeEvents("alias-a", []string{"d-1"}); err != nil {
		t.Fatalf("ack send failed: %v", err)
	}
	select {
	case ids := <-acks:
		if len(ids) != 1 || ids[0] != "d-1" {
			t.Fatalf("ack frame mismatch: %v", ids)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("ack frame never reached the gateway")
	}

	// 心跳帧刷新会话健康。
	if err := conn.WriteJSON(map[string]any{"type": "heartbeat", "version": "2026.10.0"}); err != nil {
		t.Fatal(err)
	}
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		store.mu.Lock()
		heartbeats := store.heartbeats
		store.mu.Unlock()
		if heartbeats > 0 {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatal("heartbeat was not persisted")
}

func TestGatewayChannelUnsubscribedEventBatchIsNacked(t *testing.T) {
	store := newMemoryChannelStore(newTestGateway(6, "gw-nack", "test-access-key-0006"))
	server := (&gatewayChannelTestServer{store: store, URL: startChannelServer(t, store)})

	nacks := make(chan string, 4)
	ready := make(chan struct{})
	var conn *websocket.Conn
	go func() {
		header := map[string][]string{"Authorization": {"Bearer test-access-key-0006"}}
		c, _, err := websocket.DefaultDialer.Dial(serverWSURL(server.URL)+"/v1/agent", header)
		if err != nil {
			t.Errorf("dial failed: %v", err)
			return
		}
		conn = c
		go func() {
			for {
				var frame struct {
					Type  string `json:"type"`
					Alias string `json:"alias"`
				}
				if err := c.ReadJSON(&frame); err != nil {
					return
				}
				if frame.Type == "event_nack" {
					nacks <- frame.Alias
				}
			}
		}()
		close(ready)
	}()
	<-ready
	defer func() { _ = conn.Close() }()
	<-waitForSession(t, 6)

	if err := conn.WriteJSON(map[string]any{"type": "event", "event": eventBatchFor("ghost", "d-9")}); err != nil {
		t.Fatal(err)
	}
	select {
	case alias := <-nacks:
		if alias != "ghost" {
			t.Fatalf("nack alias mismatch: %s", alias)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("unsubscribed batch was silently consumed instead of nacked")
	}
}

func TestGatewayChannelTaskPayloadRoundTrip(t *testing.T) {
	store := newMemoryChannelStore(newTestGateway(7, "gw-payload", "test-access-key-0007"))
	server := (&gatewayChannelTestServer{store: store, URL: startChannelServer(t, store)})
	SetGatewayChannelStore(store)

	received := make(chan json.RawMessage, 1)
	fake := connectFakeGateway(t, server.URL, "test-access-key-0007", func(task taskFrame) *TaskResult {
		received <- task.Payload
		return &TaskResult{Status: 204}
	})
	_ = fake
	<-waitForSession(t, 7)
	status, _, err := gatewayCall(context.Background(), newTestGateway(7, "gw-payload", "unused"), "POST", "/v1/browsers", map[string]any{"alias": "a", "nested": map[string]int{"seed": 7}}, 5*time.Second)
	if err != nil || status != 204 {
		fake.mu.Lock()
		t.Logf("diagnose: taskIDs=%d results=%#v", len(fake.taskIDs), fake.results)
		fake.mu.Unlock()
		t.Fatalf("payload task failed: status=%d err=%v", status, err)
	}
	var payload map[string]any
	if err := json.Unmarshal(<-received, &payload); err != nil {
		t.Fatalf("payload must stay valid JSON: %v", err)
	}
	if payload["alias"] != "a" {
		t.Fatalf("payload content mismatch: %v", payload)
	}
}