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

package controlplane

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"sync"
	"testing"
	"time"

	_ "modernc.org/sqlite"
)

func TestAccountStoreIsolatesAccountsAndAppliesIdempotentBatches(t *testing.T) {
	ctx := context.Background()
	root := t.TempDir()
	manager, err := OpenAccountStoreManager(root)
	if err != nil {
		t.Fatal(err)
	}
	defer manager.Close()

	scopesA := []ReportingScope{
		{ChatID: "chat-a", DataType: "read", ConfigVersion: 1},
	}
	accountA, err := manager.RegisterAccount(ctx, AccountRegistration{
		AccountID: "account-a", StableIdentity: "wechat-a", SourceNodeID: "node-a", SourceGeneration: "generation-a", Verified: true, AuthorizationVersion: 1, ReportingScopes: scopesA,
	})
	if err != nil {
		t.Fatal(err)
	}
	defer accountA.Close()
	accountB, err := manager.RegisterAccount(ctx, AccountRegistration{
		AccountID: "account-b", StableIdentity: "wechat-b", SourceNodeID: "node-b", SourceGeneration: "generation-b", Verified: true, AuthorizationVersion: 1,
		ReportingScopes: []ReportingScope{{ChatID: "chat-b", DataType: "read", ConfigVersion: 1}},
	})
	if err != nil {
		t.Fatal(err)
	}
	defer accountB.Close()
	if accountA.Path() == accountB.Path() {
		t.Fatal("accounts share a shard path")
	}
	if _, err := os.Stat(filepath.Join(root, "catalog.sqlite")); err != nil {
		t.Fatal(err)
	}

	now := time.Now().UTC().Truncate(time.Microsecond)
	batch := IngestBatch{
		BatchID: "batch-1", SourceGeneration: "generation-a", StreamKey: "messages", Sequence: 1,
		CursorStart: "0", CursorEnd: "1", PayloadHash: "hash-1", CoverageState: "complete",
		Conversations: []ConversationRecord{{ChatID: "chat-a", ChatType: "private", Title: "A", Source: "db", ObservedAt: now, DirectoryState: "active"}},
		Messages:      []MessageRecord{{MessageID: "message-1", ChatID: "chat-a", SourceMessageID: "source-1", Direction: "incoming", MessageType: "text", Text: "hello", SourceTime: now, ObservedAt: now, SourceVersion: "wx-1", PayloadHash: "message-hash"}},
	}
	result, err := accountA.ApplyBatch(ctx, batch)
	if err != nil {
		t.Fatal(err)
	}
	if result.Duplicate || result.ConfirmedSequence != 1 {
		t.Fatalf("unexpected first apply result: %+v", result)
	}
	duplicate, err := accountA.ApplyBatch(ctx, batch)
	if err != nil {
		t.Fatal(err)
	}
	if !duplicate.Duplicate || duplicate.ConfirmedSequence != 1 {
		t.Fatalf("unexpected duplicate result: %+v", duplicate)
	}
	conflicting := batch
	conflicting.PayloadHash = "different"
	if _, err := accountA.ApplyBatch(ctx, conflicting); !errors.Is(err, ErrBatchConflict) {
		t.Fatalf("expected batch conflict, got %v", err)
	}
	gap := batch
	gap.BatchID = "batch-3"
	gap.Sequence = 3
	gap.PayloadHash = "hash-3"
	if _, err := accountA.ApplyBatch(ctx, gap); !errors.Is(err, ErrBatchSequenceGap) {
		t.Fatalf("expected sequence gap, got %v", err)
	}
	unauthorized := batch
	unauthorized.BatchID = "batch-2"
	unauthorized.Sequence = 2
	unauthorized.PayloadHash = "hash-2"
	unauthorized.Messages = []MessageRecord{{MessageID: "message-b", ChatID: "chat-b", SourceMessageID: "source-b", Direction: "incoming", MessageType: "text", Text: "no", SourceTime: now, ObservedAt: now, SourceVersion: "wx-1", PayloadHash: "message-b-hash"}}
	if _, err := accountA.ApplyBatch(ctx, unauthorized); !errors.Is(err, ErrAccountNotAuthorized) {
		t.Fatalf("expected unauthorized batch, got %v", err)
	}

	messages, err := accountA.QueryMessages(ctx, "chat-a", 20, 0)
	if err != nil || len(messages) != 1 || messages[0].Text != "hello" {
		t.Fatalf("unexpected stored messages: %v %+v", err, messages)
	}
	conversations, err := accountA.QueryConversations(ctx, 20, 0)
	if err != nil || len(conversations) != 1 || conversations[0].ChatID != "chat-a" {
		t.Fatalf("unexpected stored conversations: %v %+v", err, conversations)
	}
}

func TestAccountStoreBackupRestoreAndSchemaFailureAreSafe(t *testing.T) {
	ctx := context.Background()
	root := t.TempDir()
	manager, err := OpenAccountStoreManager(root)
	if err != nil {
		t.Fatal(err)
	}
	store, err := manager.RegisterAccount(ctx, AccountRegistration{
		AccountID: "account-a", StableIdentity: "wechat-a", SourceNodeID: "node-a", SourceGeneration: "generation-a", Verified: true,
		ReportingScopes: []ReportingScope{{ChatID: "chat-a", DataType: "messages", ConfigVersion: 1}},
	})
	if err != nil {
		t.Fatal(err)
	}
	now := time.Now().UTC()
	batch1 := IngestBatch{
		BatchID: "batch-1", SourceGeneration: "generation-a", StreamKey: "messages", Sequence: 1, CursorStart: "0", CursorEnd: "1", PayloadHash: "hash-1", CoverageState: "complete",
		Messages: []MessageRecord{{MessageID: "message-1", ChatID: "chat-a", SourceMessageID: "source-1", Direction: "incoming", MessageType: "text", Text: "one", SourceTime: now, ObservedAt: now, SourceVersion: "wx", PayloadHash: "message-hash-1"}},
	}
	if _, err := store.ApplyBatch(ctx, batch1); err != nil {
		t.Fatal(err)
	}
	backup := filepath.Join(root, "backups", "account-a.sqlite")
	if err := store.Backup(ctx, backup); err != nil {
		t.Fatal(err)
	}
	if err := store.Close(); err != nil {
		t.Fatal(err)
	}
	if err := manager.RestoreAccount(ctx, "account-a", backup); err != nil {
		t.Fatal(err)
	}
	restored, err := manager.OpenAccount(ctx, "account-a")
	if err != nil {
		t.Fatal(err)
	}
	batch2 := batch1
	batch2.BatchID = "batch-2"
	batch2.Sequence = 2
	batch2.CursorStart = "1"
	batch2.CursorEnd = "2"
	batch2.PayloadHash = "hash-2"
	batch2.Messages = []MessageRecord{{MessageID: "message-2", ChatID: "chat-a", SourceMessageID: "source-2", Direction: "incoming", MessageType: "text", Text: "two", SourceTime: now.Add(time.Second), ObservedAt: now.Add(time.Second), SourceVersion: "wx", PayloadHash: "message-hash-2"}}
	if result, err := restored.ApplyBatch(ctx, batch2); err != nil || result.ConfirmedSequence != 2 {
		t.Fatalf("cursor did not continue after restore: %+v %v", result, err)
	}
	if err := restored.Close(); err != nil {
		t.Fatal(err)
	}

	corruptBackup := filepath.Join(root, "backups", "corrupt.sqlite")
	if err := copyFile(backup, corruptBackup); err != nil {
		t.Fatal(err)
	}
	corruptDB, err := sql.Open("sqlite", corruptBackup)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := corruptDB.Exec(`UPDATE schema_meta SET value = '999' WHERE key = 'schema_version'`); err != nil {
		t.Fatal(err)
	}
	if err := corruptDB.Close(); err != nil {
		t.Fatal(err)
	}
	if err := manager.RestoreAccount(ctx, "account-a", corruptBackup); err == nil {
		t.Fatal("expected unsupported schema version")
	}
	if _, err := manager.OpenAccount(ctx, "account-a"); err != nil {
		t.Fatalf("original shard should remain usable after failed restore: %v", err)
	}
}

func TestAccountStoreMaintenanceCapacityAndRevocation(t *testing.T) {
	ctx := context.Background()
	root := t.TempDir()
	manager, err := OpenAccountStoreManager(root, AccountStoreManagerOptions{Retention: time.Hour, MaxBatchBytes: 1024})
	if err != nil {
		t.Fatal(err)
	}
	defer manager.Close()
	store, err := manager.RegisterAccount(ctx, AccountRegistration{
		AccountID: "account-a", StableIdentity: "wechat-a", SourceNodeID: "node-a", SourceGeneration: "generation-a", Verified: true,
		ReportingScopes: []ReportingScope{{ChatID: "chat-a", DataType: "read", ConfigVersion: 1}},
	})
	if err != nil {
		t.Fatal(err)
	}
	old := time.Now().UTC().Add(-2 * time.Hour)
	batch := IngestBatch{
		BatchID: "old-batch", SourceGeneration: "generation-a", StreamKey: "messages", Sequence: 1, PayloadHash: "old-hash",
		Conversations: []ConversationRecord{{ChatID: "chat-a", ChatType: "private", Title: "old", Source: "db", ObservedAt: old, DirectoryState: "active"}},
		Messages:      []MessageRecord{{MessageID: "old-message", ChatID: "chat-a", SourceMessageID: "old-source", Direction: "incoming", MessageType: "text", Text: "old", SourceTime: old, ObservedAt: old, SourceVersion: "wx", PayloadHash: "old-message-hash"}},
	}
	if _, err := store.ApplyBatch(ctx, batch); err != nil {
		t.Fatal(err)
	}
	report, err := manager.RunMaintenance(ctx, time.Now().UTC())
	if err != nil {
		t.Fatal(err)
	}
	if report.DeletedMessages != 1 || report.DeletedConversations != 1 || report.Checkpointed != 1 {
		t.Fatalf("unexpected maintenance report: %+v", report)
	}
	if _, err := store.QueryMessages(ctx, "chat-a", 20, 0); err != nil {
		t.Fatal(err)
	}
	if err := manager.SetScopes(ctx, "account-a", nil); err != nil {
		t.Fatal(err)
	}
	if _, err := store.QueryMessages(ctx, "chat-a", 20, 0); !errors.Is(err, ErrAccountNotAuthorized) {
		t.Fatalf("expected revoked query to be rejected, got %v", err)
	}
	if err := manager.SetScopes(ctx, "account-a", []ReportingScope{{ChatID: "chat-a", DataType: "read", ConfigVersion: 2}}); err != nil {
		t.Fatal(err)
	}

	oversized := batch
	oversized.BatchID = "oversized"
	oversized.Sequence = 2
	oversized.PayloadHash = "oversized-hash"
	oversized.Messages = []MessageRecord{{MessageID: "oversized-message", ChatID: "chat-a", SourceMessageID: "oversized-source", Direction: "incoming", MessageType: "text", Text: strings.Repeat("x", 2000), SourceTime: time.Now().UTC(), ObservedAt: time.Now().UTC(), SourceVersion: "wx", PayloadHash: "oversized-message-hash"}}
	if _, err := store.ApplyBatch(ctx, oversized); !errors.Is(err, ErrAccountCapacityExceeded) {
		t.Fatalf("expected oversized batch to be rejected by capacity budget, got %v", err)
	}
}

func TestAccountStoreSyntheticConcurrentReadWrite(t *testing.T) {
	ctx := context.Background()
	root := t.TempDir()
	manager, err := OpenAccountStoreManager(root, AccountStoreManagerOptions{MaxShardBytes: 8 * 1024 * 1024, MaxBatchBytes: 256 * 1024})
	if err != nil {
		t.Fatal(err)
	}
	defer manager.Close()
	const accountCount = 8
	const batchesPerAccount = 5
	const messagesPerBatch = 20
	for accountIndex := 0; accountIndex < accountCount; accountIndex++ {
		accountID := fmt.Sprintf("account-%02d", accountIndex)
		chatID := fmt.Sprintf("chat-%02d", accountIndex)
		store, err := manager.RegisterAccount(ctx, AccountRegistration{
			AccountID: accountID, StableIdentity: accountID, SourceNodeID: "pressure-node", SourceGeneration: "pressure-generation", Verified: true,
			ReportingScopes: []ReportingScope{{ChatID: chatID, DataType: "read", ConfigVersion: 1}},
		})
		if err != nil {
			t.Fatal(err)
		}
		for batchIndex := 0; batchIndex < batchesPerAccount; batchIndex++ {
			messages := make([]MessageRecord, 0, messagesPerBatch)
			for messageIndex := 0; messageIndex < messagesPerBatch; messageIndex++ {
				messages = append(messages, MessageRecord{MessageID: fmt.Sprintf("%s-%d-%d", accountID, batchIndex, messageIndex), ChatID: chatID, SourceMessageID: fmt.Sprintf("source-%d-%d", batchIndex, messageIndex), Direction: "incoming", MessageType: "text", Text: "synthetic", SourceTime: time.Now().UTC(), ObservedAt: time.Now().UTC(), SourceVersion: "pressure", PayloadHash: fmt.Sprintf("hash-%d-%d", batchIndex, messageIndex)})
			}
			if _, err := store.ApplyBatch(ctx, IngestBatch{BatchID: fmt.Sprintf("%s-batch-%d", accountID, batchIndex), SourceGeneration: "pressure-generation", StreamKey: "messages", Sequence: int64(batchIndex + 1), PayloadHash: fmt.Sprintf("batch-hash-%s-%d", accountID, batchIndex), Messages: messages}); err != nil {
				t.Fatal(err)
			}
		}
		if err := store.Close(); err != nil {
			t.Fatal(err)
		}
	}
	start := time.Now()
	var wg sync.WaitGroup
	var mu sync.Mutex
	var slowest time.Duration
	var rowsRead int
	for accountIndex := 0; accountIndex < accountCount; accountIndex++ {
		accountID := fmt.Sprintf("account-%02d", accountIndex)
		chatID := fmt.Sprintf("chat-%02d", accountIndex)
		wg.Add(1)
		go func() {
			defer wg.Done()
			store, err := manager.OpenAccount(ctx, accountID)
			if err != nil {
				t.Errorf("open %s: %v", accountID, err)
				return
			}
			defer store.Close()
			localStart := time.Now()
			for i := 0; i < 20; i++ {
				items, err := store.QueryMessages(ctx, chatID, 100, 0)
				if err != nil {
					t.Errorf("query %s: %v", accountID, err)
					return
				}
				mu.Lock()
				rowsRead += len(items)
				mu.Unlock()
			}
			mu.Lock()
			if elapsed := time.Since(localStart); elapsed > slowest {
				slowest = elapsed
			}
			mu.Unlock()
		}()
	}
	wg.Wait()
	t.Logf("synthetic accounts=%d messages=%d reads=%d elapsed=%s slowest-account=%s", accountCount, accountCount*batchesPerAccount*messagesPerBatch, rowsRead, time.Since(start), slowest)
	if rowsRead != accountCount*batchesPerAccount*messagesPerBatch*20 {
		t.Fatalf("unexpected synthetic read count: %d", rowsRead)
	}
}

func TestAccountStoreAddsIndexesToExistingShard(t *testing.T) {
	ctx := context.Background()
	root := t.TempDir()
	manager, err := OpenAccountStoreManager(root)
	if err != nil {
		t.Fatal(err)
	}
	store, err := manager.RegisterAccount(ctx, AccountRegistration{AccountID: "account-a", StableIdentity: "wechat-a", SourceNodeID: "node-a", SourceGeneration: "generation-a", Verified: true, ReportingScopes: []ReportingScope{{ChatID: "chat-a", DataType: "read", ConfigVersion: 1}}})
	if err != nil {
		t.Fatal(err)
	}
	if _, err := store.db.ExecContext(ctx, `DROP INDEX idx_messages_observed_at; DROP INDEX idx_conversations_activity;`); err != nil {
		t.Fatal(err)
	}
	if err := store.Close(); err != nil {
		t.Fatal(err)
	}
	if err := manager.Close(); err != nil {
		t.Fatal(err)
	}
	manager, err = OpenAccountStoreManager(root)
	if err != nil {
		t.Fatal(err)
	}
	defer manager.Close()
	store, err = manager.OpenAccount(ctx, "account-a")
	if err != nil {
		t.Fatal(err)
	}
	defer store.Close()
	rows, err := store.db.QueryContext(ctx, `SELECT name FROM sqlite_master WHERE type = 'index' AND name IN ('idx_messages_observed_at', 'idx_conversations_activity') ORDER BY name`)
	if err != nil {
		t.Fatal(err)
	}
	defer rows.Close()
	var names []string
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			t.Fatal(err)
		}
		names = append(names, name)
	}
	if err := rows.Err(); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(names, []string{"idx_conversations_activity", "idx_messages_observed_at"}) {
		t.Fatalf("indexes were not restored on existing shard: %v", names)
	}
}

func TestAccountStoreRejectsExpiredAuthorization(t *testing.T) {
	ctx := context.Background()
	manager, err := OpenAccountStoreManager(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	defer manager.Close()
	expired := time.Now().UTC().Add(-time.Minute)
	store, err := manager.RegisterAccount(ctx, AccountRegistration{
		AccountID: "account-a", StableIdentity: "wechat-a", SourceNodeID: "node-a", SourceGeneration: "generation-a", Verified: true,
		AuthorizationExpiresAt: &expired, ReportingScopes: []ReportingScope{{ChatID: "chat-a", DataType: "messages", ConfigVersion: 1}},
	})
	if err != nil {
		t.Fatal(err)
	}
	defer store.Close()
	_, err = store.ApplyBatch(ctx, IngestBatch{BatchID: "batch-1", SourceGeneration: "generation-a", StreamKey: "messages", Sequence: 1, PayloadHash: "hash"})
	if !errors.Is(err, ErrAccountNotAuthorized) {
		t.Fatalf("expected expired authorization error, got %v", err)
	}
}

func copyFile(source, target string) error {
	data, err := os.ReadFile(source)
	if err != nil {
		return err
	}
	return os.WriteFile(target, data, 0o600)
}