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

package api

// 网关出站通道测试基建：内存桩 store + 平台侧 WS 会话服务器 + 假网关客户端。

import (
	"context"
	"encoding/json"
	"net"
	"net/http"
	"strings"
	"sync"
	"testing"
	"time"

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

type memoryChannelStore struct {
	mu          sync.Mutex
	gateways    map[string]hub.Gateway // access_key -> gateway
	tasks       map[int64]*hub.GatewayTask
	taskSeq     int64
	heartbeats  int
	onlineCalls int
	interrupted int
}

func newMemoryChannelStore(gateways ...hub.Gateway) *memoryChannelStore {
	store := &memoryChannelStore{gateways: map[string]hub.Gateway{}, tasks: map[int64]*hub.GatewayTask{}}
	for _, gateway := range gateways {
		store.gateways[gateway.AccessKey] = gateway
	}
	return store
}

func (s *memoryChannelStore) TouchGatewaySession(_ context.Context, gatewayID int64, version string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.heartbeats++
	for key, gateway := range s.gateways {
		if gateway.ID == gatewayID {
			gateway.Version = version
			gateway.Online = true
			now := time.Now()
			gateway.LastSeenAt = &now
			s.gateways[key] = gateway
		}
	}
	return nil
}

func (s *memoryChannelStore) SetGatewayOnline(_ context.Context, gatewayID int64, online bool) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.onlineCalls++
	for key, gateway := range s.gateways {
		if gateway.ID == gatewayID {
			gateway.Online = online
			s.gateways[key] = gateway
		}
	}
	return nil
}

func (s *memoryChannelStore) InterruptDispatchedTasks(_ context.Context, _ int64, _ string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.interrupted++
	return nil
}

func (s *memoryChannelStore) GetGatewayByAccessKey(_ context.Context, accessKey string) (hub.Gateway, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	accessKey = strings.TrimPrefix(accessKey, "Bearer ")
	if gateway, ok := s.gateways[accessKey]; ok {
		return gateway, nil
	}
	return hub.Gateway{}, hub.ErrNotFound
}

func (s *memoryChannelStore) CreateGatewayTask(_ context.Context, gatewayID int64, method, path string, payload []byte, timeout time.Duration) (hub.GatewayTask, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.taskSeq++
	timeoutSeconds := int(timeout / time.Second)
	if timeoutSeconds < 1 {
		timeoutSeconds = 1
	}
	task := hub.GatewayTask{ID: s.taskSeq, GatewayID: gatewayID, Method: method, Path: path,
		Payload: payload, Status: hub.TaskPending, TimeoutSeconds: timeoutSeconds}
	s.tasks[task.ID] = &task
	return task, nil
}

func (s *memoryChannelStore) DispatchGatewayTask(_ context.Context, taskID int64) (bool, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if task, ok := s.tasks[taskID]; ok && task.Status == hub.TaskPending {
		task.Status = hub.TaskDispatched
		return true, nil
	}
	return false, nil
}

func (s *memoryChannelStore) FinishGatewayTask(_ context.Context, taskID int64, status string, responseStatus int, taskError string, responseBytes int64) (bool, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if task, ok := s.tasks[taskID]; ok && (task.Status == hub.TaskDispatched || task.Status == hub.TaskPending) {
		task.Status = status
		task.ResponseStatus = responseStatus
		task.Error = taskError
		task.ResponseBytes = responseBytes
		return true, nil
	}
	return false, nil
}

func (s *memoryChannelStore) NextPendingGatewayTasks(_ context.Context, gatewayID int64, limit int) ([]hub.GatewayTask, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	tasks := []hub.GatewayTask{}
	for id := int64(1); id <= s.taskSeq && len(tasks) < limit; id++ {
		if task, ok := s.tasks[id]; ok && task.GatewayID == gatewayID && task.Status == hub.TaskPending {
			tasks = append(tasks, *task)
		}
	}
	return tasks, nil
}

func (s *memoryChannelStore) ResetGatewayConnections(_ context.Context) error { return nil }

func (s *memoryChannelStore) CreateTaskRow(gatewayID int64, method, path string, payload []byte) hub.GatewayTask {
	task, _ := s.CreateGatewayTask(context.Background(), gatewayID, method, path, payload, 30*time.Second)
	return task
}

// startChannelServer 启动与生产一致的 fiber WS 服务（真实 listener + FastHTTPUpgrader）。
func startChannelServer(t *testing.T, store gatewaySessionStore) string {
	t.Helper()
	app := fiber.New()
	app.Get("/v1/agent", func(c fiber.Ctx) error {
		authorization := string(c.RequestCtx().Request.Header.Peek("Authorization"))
		upgrader := websocket.FastHTTPUpgrader{}
		return upgrader.Upgrade(c.RequestCtx(), func(conn *websocket.Conn) {
			serveGatewaySession(conn, authorization, store)
		})
	})
	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("channel listener: %v", err)
	}
	go func() { _ = app.Listener(listener, fiber.ListenConfig{DisableStartupMessage: true}) }()
	t.Cleanup(func() { _ = listener.Close() })
	return "http://" + listener.Addr().String()
}

// gatewayChannelTestServer 旧签名兼容：URL 指向真实 fiber listener。
type gatewayChannelTestServer struct {
	store *memoryChannelStore
	URL   string
}

type fakeChannelGateway struct {
	mu      sync.Mutex
	conn    *websocket.Conn
	taskIDs chan int64
	results map[int64]TaskResult // 已发送的 result，供测试诊断
}

// WriteJSON 串行化对连接的写，避免与任务 result/心跳并发写帧。
func (g *fakeChannelGateway) WriteJSON(value any) error {
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.conn.WriteJSON(value)
}

// connectFakeGateway 模拟网关出站连接：收到 task 帧后用 handler 生成 result 帧。
func connectFakeGateway(t *testing.T, serverURL, accessKey string, handle func(task taskFrame) *TaskResult) *fakeChannelGateway {
	t.Helper()
	header := http.Header{"Authorization": []string{"Bearer " + accessKey}}
	conn, _, err := websocket.DefaultDialer.Dial(serverWSURL(serverURL)+"/v1/agent", header)
	if err != nil {
		t.Fatalf("fake gateway dial failed: %v", err)
	}
	gateway := &fakeChannelGateway{conn: conn, taskIDs: make(chan int64, 8), results: map[int64]TaskResult{}}
	go func() {
		for {
			var raw struct {
				Type    string          `json:"type"`
				ID      int64           `json:"id"`
				Method  string          `json:"method"`
				Path    string          `json:"path"`
				Payload json.RawMessage `json:"payload"`
			}
			if err := conn.ReadJSON(&raw); err != nil {
				return
			}
			if raw.Type != "task" {
				continue
			}
			select {
			case gateway.taskIDs <- raw.ID:
			default:
			}
			if result := handle(taskFrame{ID: raw.ID, Method: raw.Method, Path: raw.Path, Payload: raw.Payload}); result != nil {
				gateway.mu.Lock()
				writeErr := conn.WriteJSON(resultFrame{Type: "result", ID: raw.ID, TaskResult: *result})
				gateway.results[raw.ID] = *result
				gateway.mu.Unlock()
				if writeErr != nil {
					t.Logf("fake gateway result write failed for task %d: %v", raw.ID, writeErr)
				}
			}
		}
	}()
	// 平台 60s 读超时依赖心跳；假网关必须周期回心跳保持会话。
	tickerStop := make(chan struct{})
	go func() {
		ticker := time.NewTicker(20 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				gateway.mu.Lock()
				_ = conn.WriteJSON(map[string]any{"type": "heartbeat", "version": "test"})
				gateway.mu.Unlock()
			case <-tickerStop:
				return
			}
		}
	}()
	t.Cleanup(func() { close(tickerStop) })
	t.Cleanup(func() { _ = conn.Close() })
	return gateway
}

func newTestGateway(id int64, name, accessKey string) hub.Gateway {
	return hub.Gateway{ID: id, Name: name, AccessKey: accessKey}
}

func eventBatchFor(alias, deliveryID string) EventBatch {
	return EventBatch{Subscription: "sub-1", Alias: alias, Deliveries: []creator.ListenerDelivery{{
		DeliveryID: deliveryID, Kind: "open",
	}}}
}

// serverWSURL 把 http:// 地址转为 ws:// 地址。
func serverWSURL(httpURL string) string {
	return "ws://" + strings.TrimPrefix(httpURL, "http://")
}