unibus/pkg/membership/migrate_test.go

package membership

import (
	"path/filepath"
	"reflect"
	"sort"
	"testing"
	"time"

	"github.com/enmanuel/unibus/pkg/embeddednats"
	"github.com/nats-io/nats.go"
	"github.com/nats-io/nats.go/jetstream"
)

// seedSQLite populates a SQLite store with a representative control plane: two
// rooms (one rekeyed to epoch 2 with a removed member's keys left behind), a few
// members and sealed keys, and a user allowlist with one revoked entry. It
// returns the populated *sqliteStore and its file path.
func seedSQLite(t *testing.T) (*sqliteStore, string) {
	t.Helper()
	path := filepath.Join(t.TempDir(), "seed.db")
	s, err := openSQLite(path)
	if err != nil {
		t.Fatalf("openSQLite: %v", err)
	}

	r1 := RoomInfo{RoomID: newULID(), Subject: "room.alpha", Encrypt: true, Persist: true, SignMsgs: true, OwnerEndpoint: "ep-owner1"}
	if err := s.CreateRoom(r1, []byte("o1-sign"), []byte("o1-kex"), []byte("o1-sealed-e1")); err != nil {
		t.Fatalf("create r1: %v", err)
	}
	if err := s.AddMember(r1.RoomID, Member{Endpoint: "ep-bob", Role: "member", SignPub: []byte("bob-sign"), KexPub: []byte("bob-kex")}, 1, []byte("bob-sealed-e1")); err != nil {
		t.Fatalf("add bob: %v", err)
	}
	// Rekey r1 to epoch 2 (owner keeps a key at the new epoch).
	if err := s.BumpEpoch(r1.RoomID, 2); err != nil {
		t.Fatalf("bump: %v", err)
	}
	if err := s.PutSealedKeys(r1.RoomID, 2, map[string][]byte{"ep-owner1": []byte("o1-sealed-e2")}); err != nil {
		t.Fatalf("put keys e2: %v", err)
	}

	r2 := RoomInfo{RoomID: newULID(), Subject: "room.beta", Encrypt: false, Persist: false, SignMsgs: false, OwnerEndpoint: "ep-owner2"}
	if err := s.CreateRoom(r2, []byte("o2-sign"), []byte("o2-kex"), nil); err != nil {
		t.Fatalf("create r2: %v", err)
	}

	if err := s.AddUser("aa11", "alice", RoleAdmin); err != nil {
		t.Fatalf("add alice: %v", err)
	}
	if err := s.AddUser("bb22", "bob", RoleMember); err != nil {
		t.Fatalf("add bob user: %v", err)
	}
	if err := s.AddUser("cc33", "carol", RoleMember); err != nil {
		t.Fatalf("add carol: %v", err)
	}
	if err := s.RevokeUser("cc33"); err != nil {
		t.Fatalf("revoke carol: %v", err)
	}
	return s, path
}

// normalizeSnapshot sorts every slice in a Snapshot so two snapshots from
// different backends can be compared regardless of enumeration order.
func normalizeSnapshot(snap *Snapshot) {
	sort.Slice(snap.Rooms, func(i, j int) bool { return snap.Rooms[i].RoomID < snap.Rooms[j].RoomID })
	for _, ms := range snap.Members {
		sort.Slice(ms, func(i, j int) bool { return ms[i].Endpoint < ms[j].Endpoint })
	}
	sort.Slice(snap.Keys, func(i, j int) bool {
		a, b := snap.Keys[i], snap.Keys[j]
		if a.RoomID != b.RoomID {
			return a.RoomID < b.RoomID
		}
		if a.Endpoint != b.Endpoint {
			return a.Endpoint < b.Endpoint
		}
		return a.Epoch < b.Epoch
	})
	sort.Slice(snap.Users, func(i, j int) bool { return snap.Users[i].SignPub < snap.Users[j].SignPub })
}

func newJS(t *testing.T) jetstream.JetStream {
	t.Helper()
	ns, err := embeddednats.StartServer(embeddednats.ServerConfig{
		StoreDir: t.TempDir(),
		Host:     "127.0.0.1",
		Port:     kvFreePort(t),
	})
	if err != nil {
		t.Fatalf("embedded nats: %v", err)
	}
	nc, err := nats.Connect(ns.ClientURL())
	if err != nil {
		ns.Shutdown()
		t.Fatalf("nats connect: %v", err)
	}
	js, err := jetstream.New(nc)
	if err != nil {
		nc.Close()
		ns.Shutdown()
		t.Fatalf("jetstream: %v", err)
	}
	t.Cleanup(func() { nc.Close(); ns.Shutdown(); ns.WaitForShutdown() })
	return js
}

// TestMigrateSQLiteToKVParity is the parity test the issue mandates: after the
// migration, the KV store holds exactly the SQLite source's state.
func TestMigrateSQLiteToKVParity(t *testing.T) {
	src, path := seedSQLite(t)
	srcSnap, err := src.ExportSnapshot()
	if err != nil {
		t.Fatalf("export sqlite: %v", err)
	}
	src.Close() // release the file before the migration reopens it

	js := newJS(t)
	report, err := MigrateSQLiteToKV(path, js, JetStreamConfig{Replicas: 1, OpTimeout: 5 * time.Second})
	if err != nil {
		t.Fatalf("migrate: %v", err)
	}
	if report.Rooms != 2 || report.Users != 3 {
		t.Fatalf("report mismatch: %+v", report)
	}

	kv, err := OpenJetStream(js, JetStreamConfig{Replicas: 1, OpTimeout: 5 * time.Second})
	if err != nil {
		t.Fatalf("open kv: %v", err)
	}
	kvSnap, err := kv.(*jetstreamStore).ExportSnapshot()
	if err != nil {
		t.Fatalf("export kv: %v", err)
	}

	normalizeSnapshot(srcSnap)
	normalizeSnapshot(kvSnap)
	if !reflect.DeepEqual(srcSnap, kvSnap) {
		t.Fatalf("parity mismatch after migration:\n sqlite=%+v\n kv=    %+v", srcSnap, kvSnap)
	}
}

// TestMigrateSQLiteToKVIdempotent: running the migration twice converges to the
// same KV state (every write is an overwrite). A second run must not duplicate
// or corrupt anything.
func TestMigrateSQLiteToKVIdempotent(t *testing.T) {
	src, path := seedSQLite(t)
	srcSnap, _ := src.ExportSnapshot()
	src.Close()

	js := newJS(t)
	if _, err := MigrateSQLiteToKV(path, js, JetStreamConfig{Replicas: 1}); err != nil {
		t.Fatalf("migrate run 1: %v", err)
	}
	if _, err := MigrateSQLiteToKV(path, js, JetStreamConfig{Replicas: 1}); err != nil {
		t.Fatalf("migrate run 2: %v", err)
	}

	kv, _ := OpenJetStream(js, JetStreamConfig{Replicas: 1})
	kvSnap, err := kv.(*jetstreamStore).ExportSnapshot()
	if err != nil {
		t.Fatalf("export kv: %v", err)
	}
	normalizeSnapshot(srcSnap)
	normalizeSnapshot(kvSnap)
	if !reflect.DeepEqual(srcSnap, kvSnap) {
		t.Fatalf("idempotency broken: a second migration changed the KV state\n sqlite=%+v\n kv=    %+v", srcSnap, kvSnap)
	}
}

// TestBackupSQLiteCreatesConsistentCopy verifies the pre-migration backup is a
// real, openable copy holding the same data.
func TestBackupSQLiteCreatesConsistentCopy(t *testing.T) {
	src, path := seedSQLite(t)
	srcSnap, _ := src.ExportSnapshot()
	src.Close()

	bak, err := BackupSQLite(path)
	if err != nil {
		t.Fatalf("backup: %v", err)
	}
	restored, err := openSQLite(bak)
	if err != nil {
		t.Fatalf("open backup: %v", err)
	}
	defer restored.Close()
	bakSnap, err := restored.ExportSnapshot()
	if err != nil {
		t.Fatalf("export backup: %v", err)
	}
	normalizeSnapshot(srcSnap)
	normalizeSnapshot(bakSnap)
	if !reflect.DeepEqual(srcSnap, bakSnap) {
		t.Fatalf("backup is not a faithful copy")
	}
}