Agent-System-Backend/internal/server/recovery_multi_instance_test.go

package server

import (
	"database/sql"
	"encoding/json"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"agent-system-backend/internal/store"
)

func twinRecoveryStores(t *testing.T) (*store.Store, *store.Store) {
	t.Helper()
	dsn := os.Getenv("AGENT_PG_TEST_DSN")
	if dsn == "" {
		path := filepath.Join(t.TempDir(), "shared.db")
		first, err := store.OpenSQLite(path)
		if err != nil {
			t.Fatal(err)
		}
		second, err := store.OpenSQLite(path)
		if err != nil {
			first.Close()
			t.Fatal(err)
		}
		t.Cleanup(func() { second.Close(); first.Close() })
		return first, second
	}
	admin, err := sql.Open("pgx", dsn)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := admin.Exec(`CREATE EXTENSION IF NOT EXISTS pg_trgm`); err != nil {
		t.Fatal(err)
	}
	schema := "t_" + strings.ToLower(store.NewID("recovery"))
	if _, err := admin.Exec(`CREATE SCHEMA ` + schema); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _, _ = admin.Exec(`DROP SCHEMA IF EXISTS ` + schema + ` CASCADE`); admin.Close() })
	isolate := store.AppendPGRuntimeParam(dsn, "search_path", schema+",public")
	first, err := store.OpenPostgres(isolate, filepath.Join(t.TempDir(), "a-artifacts"))
	if err != nil {
		t.Fatal(err)
	}
	second, err := store.OpenPostgres(isolate, filepath.Join(t.TempDir(), "b-artifacts"))
	if err != nil {
		first.Close()
		t.Fatal(err)
	}
	t.Cleanup(func() { second.Close(); first.Close() })
	return first, second
}

func TestRecoveryAcrossTwoHTTPInstancesAndIndependentPools(t *testing.T) {
	first, second := twinRecoveryStores(t)
	t.Logf("recovery HTTP dialect=%s two independent stores", first.Dialect())
	hA := newAuthTestServerOnStore(t, first, t.TempDir()).Handler()
	hB := newAuthTestServerOnStore(t, second, t.TempDir()).Handler()
	adminCookie := bootAdmin(t, hA, first)
	account, err := first.CreatePendingAccount("admin", "two-instance-owner", store.AccountProfile{})
	if err != nil {
		t.Fatal(err)
	}
	activation, err := first.InitializePendingAccount(account.ID, account.StateVersion, account.InitializationGeneration)
	if err != nil {
		t.Fatal(err)
	}
	setup := do(hB, "POST", "/api/auth/login-key/redeem", `{"key":"`+activation+`"}`, "")
	if setup.Code != http.StatusOK {
		t.Fatalf("second instance activation redeem: %d %s", setup.Code, setup.Body.String())
	}
	if done := do(hA, "POST", "/api/auth/password", `{"next":"two-instance-initial"}`, cookieOf(setup)); done.Code != http.StatusOK {
		t.Fatalf("first instance initial set: %d %s", done.Code, done.Body.String())
	}
	login := do(hA, "POST", "/api/auth/login", `{"owner":"two-instance-owner","password":"two-instance-initial"}`, "")
	if login.Code != http.StatusOK {
		t.Fatalf("first instance login: %d %s", login.Code, login.Body.String())
	}
	ownerCookie := cookieOf(login)
	current, err := second.GetAccount(account.ID)
	if err != nil {
		t.Fatal(err)
	}
	proposalResponse := do(hA, "POST", "/api/admin/users/"+account.ID+"/recovery-proposals",
		`{"expectedStateVersion":`+jsonNumber(current.StateVersion)+`,"idempotencyKey":"two-instance-proposal"}`, adminCookie)
	if proposalResponse.Code != http.StatusOK {
		t.Fatalf("proposal: %d %s", proposalResponse.Code, proposalResponse.Body.String())
	}
	var proposal store.RecoveryProposal
	if err := json.Unmarshal(proposalResponse.Body.Bytes(), &proposal); err != nil {
		t.Fatal(err)
	}
	confirmation := do(hB, "POST", "/api/auth/recovery-proposals/"+proposal.ID+"/confirm", `{"expectedVersion":1}`, ownerCookie)
	if confirmation.Code != http.StatusOK {
		t.Fatalf("second instance confirm: %d %s", confirmation.Code, confirmation.Body.String())
	}
	var confirmed struct {
		Grant string `json:"grant"`
	}
	if err := json.Unmarshal(confirmation.Body.Bytes(), &confirmed); err != nil || confirmed.Grant == "" {
		t.Fatalf("grant response: %v", err)
	}
	redeemed := do(hA, "POST", "/api/auth/recovery-grants/redeem", `{"key":"`+confirmed.Grant+`"}`, "")
	if redeemed.Code != http.StatusOK {
		t.Fatalf("first instance grant redeem: %d %s", redeemed.Code, redeemed.Body.String())
	}
	if completed := do(hB, "POST", "/api/auth/password", `{"next":"two-instance-recovered"}`, cookieOf(redeemed)); completed.Code != http.StatusOK {
		t.Fatalf("second instance recovery completion: %d %s", completed.Code, completed.Body.String())
	}
	if old := do(hB, "GET", "/api/me", "", ownerCookie); old.Code != http.StatusUnauthorized {
		t.Fatalf("old proof session survived: %d", old.Code)
	}
	if fresh := do(hA, "POST", "/api/auth/login", `{"owner":"two-instance-owner","password":"two-instance-recovered"}`, ""); fresh.Code != http.StatusOK {
		t.Fatalf("recovered password login: %d %s", fresh.Code, fresh.Body.String())
	}
	after, err := first.GetAccount(account.ID)
	if err != nil || after.CredentialOrigin != store.CredentialOriginOwnerRecovery || after.ID != account.ID {
		t.Fatalf("persistent identity or provenance changed: %+v %v", after, err)
	}
}