wx-win-agent/control-plane/data_routes_test.go

package controlplane

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"testing"
	"time"
)

func TestDataBatchRoundTripAndWebQueriesUseAccountShard(t *testing.T) {
	server, err := NewServer(ServerConfig{
		DataFile:         filepath.Join(t.TempDir(), "control-plane.json"),
		NodeTokens:       map[string]string{"node-a": "secret-a"},
		WebUsers:         map[string]string{"admin": "web-secret"},
		HeartbeatTimeout: time.Minute,
		LeaseTTL:         time.Minute,
	})
	if err != nil {
		t.Fatal(err)
	}
	defer server.Close()
	httpServer := httptest.NewServer(server.Handler())
	defer httpServer.Close()
	client := httpServer.Client()

	register := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/nodes/register", "Bearer secret-a", NodeRegistration{
		NodeID: "node-a", ConnectionID: "connection-a", AgentVersion: "test", ProtocolVersion: ProtocolVersion,
		Capabilities: []string{"heartbeat", "poll-tasks", "sync-data"},
		Accounts:     []AccountSummary{{AccountID: "account-a", Active: true, Verified: true, AllowedChats: []AllowedChatSummary{{ChatID: "chat-a", ChatType: ChatPrivate}}}},
	})
	if register.Code != http.StatusOK {
		t.Fatalf("register status = %d: %s", register.Code, register.Body.String())
	}

	login := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/auth/login", "", map[string]string{"username": "admin", "password": "web-secret"})
	var session struct {
		AccessToken string `json:"access_token"`
	}
	decodeBody(t, login, &session)
	webAuth := "Bearer " + session.AccessToken

	now := time.Now().UTC().Truncate(time.Millisecond)
	batch := dataBatchRequest{
		NodeID: "node-a", AccountID: "account-a", BatchID: "batch-1", SourceGeneration: "account-a", StreamKey: "messages", Sequence: 1,
		CursorStart: "{}", CursorEnd: `{"chat-a\u001fmessage/a.db":1}`, PayloadHash: "batch-hash", CoverageState: "complete",
		Conversations: []dataConversationRequest{{ChatID: "chat-a", ChatType: ChatPrivate, Title: "测试会话", Source: "db", ObservedAt: now, DirectoryState: "observed"}},
		Messages:      []dataMessageRequest{{MessageID: "chat-a:local:1", ChatID: "chat-a", ChatType: ChatPrivate, SourceMessageID: "1", Direction: "incoming", MessageType: "text", Text: "hello", SourceTime: now, ObservedAt: now, SourceVersion: "test", PayloadHash: "message-hash"}},
	}
	posted := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/data/batches", "Bearer secret-a", batch)
	if posted.Code != http.StatusOK {
		t.Fatalf("batch status = %d: %s", posted.Code, posted.Body.String())
	}
	var ack dataBatchAck
	decodeBody(t, posted, &ack)
	if !ack.Accepted || ack.Duplicate {
		t.Fatalf("unexpected batch ack: %+v", ack)
	}

	duplicate := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/data/batches", "Bearer secret-a", batch)
	if duplicate.Code != http.StatusOK {
		t.Fatalf("duplicate status = %d: %s", duplicate.Code, duplicate.Body.String())
	}
	decodeBody(t, duplicate, &ack)
	if !ack.Duplicate {
		t.Fatalf("expected duplicate ack: %+v", ack)
	}

	conversations := doJSON(t, client, http.MethodGet, httpServer.URL+"/v1/data/accounts/account-a/conversations?limit=20", webAuth, nil)
	if conversations.Code != http.StatusOK {
		t.Fatalf("conversation query status = %d: %s", conversations.Code, conversations.Body.String())
	}
	var conversationBody struct {
		Items []dataConversationView `json:"items"`
		Sync  dataSyncView           `json:"sync"`
	}
	decodeBody(t, conversations, &conversationBody)
	if len(conversationBody.Items) != 1 || conversationBody.Items[0].ChatID != "chat-a" || conversationBody.Sync.State != "complete" {
		t.Fatalf("unexpected conversation body: %+v", conversationBody)
	}

	messages := doJSON(t, client, http.MethodGet, httpServer.URL+"/v1/data/accounts/account-a/messages?chat_id=chat-a&limit=20", webAuth, nil)
	if messages.Code != http.StatusOK {
		t.Fatalf("message query status = %d: %s", messages.Code, messages.Body.String())
	}
	var messageBody struct {
		Items []dataMessageView `json:"items"`
	}
	decodeBody(t, messages, &messageBody)
	if len(messageBody.Items) != 1 || messageBody.Items[0].Text != "hello" {
		t.Fatalf("unexpected message body: %+v", messageBody)
	}

	refresh := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/data/accounts/account-a/refresh", webAuth, nil)
	if refresh.Code != http.StatusAccepted {
		t.Fatalf("refresh status = %d: %s", refresh.Code, refresh.Body.String())
	}
}

func TestDataBatchCannotCrossReportingScope(t *testing.T) {
	server, err := NewServer(ServerConfig{
		DataFile: filepath.Join(t.TempDir(), "control-plane.json"), NodeTokens: map[string]string{"node-a": "secret-a"}, WebUsers: map[string]string{"admin": "web-secret"},
		HeartbeatTimeout: time.Minute, LeaseTTL: time.Minute,
	})
	if err != nil {
		t.Fatal(err)
	}
	defer server.Close()
	httpServer := httptest.NewServer(server.Handler())
	defer httpServer.Close()
	client := httpServer.Client()
	register := NodeRegistration{NodeID: "node-a", ConnectionID: "connection-a", AgentVersion: "test", ProtocolVersion: ProtocolVersion, Accounts: []AccountSummary{{AccountID: "account-a", Verified: true, AllowedChats: []AllowedChatSummary{{ChatID: "chat-a", ChatType: ChatPrivate}}}}}
	if response := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/nodes/register", "Bearer secret-a", register); response.Code != http.StatusOK {
		t.Fatalf("register: %d %s", response.Code, response.Body.String())
	}
	body := dataBatchRequest{NodeID: "node-a", AccountID: "account-a", BatchID: "denied", SourceGeneration: "account-a", StreamKey: "messages", Sequence: 1, PayloadHash: "hash", CursorStart: "{}", CursorEnd: "{}", Messages: []dataMessageRequest{{MessageID: "denied", ChatID: "chat-b", ChatType: ChatPrivate, SourceTime: time.Now().UTC(), ObservedAt: time.Now().UTC(), PayloadHash: "message"}}}
	response := doJSON(t, client, http.MethodPost, httpServer.URL+"/v1/data/batches", "Bearer secret-a", body)
	if response.Code != http.StatusForbidden {
		t.Fatalf("denied status = %d: %s", response.Code, response.Body.String())
	}
	var payload map[string]any
	if err := json.Unmarshal([]byte(response.Body.String()), &payload); err != nil {
		t.Fatal(err)
	}
}