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) } }