Agent-System/internal/authz/authorizer_pg_test.go

package authz

import (
	"context"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"

	"chat-agent-system/internal/db"
	"chat-agent-system/internal/docduty"
)

func TestUpsertProjectMemberDuty_ConcurrentLastWriter_PG(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping PostgreSQL duty concurrency in -short")
	}
	dsn := db.SkipIfNoTestDatabaseURL(t)
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()
	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		t.Fatalf("postgres unavailable: %v", err)
	}
	defer pool.Close()
	if err := pool.Ping(ctx); err != nil {
		t.Fatalf("postgres ping failed: %v", err)
	}
	if err := db.RunMigrationsOnceForTest(ctx, pool, dsn); err != nil {
		t.Fatalf("migration failed: %v", err)
	}

	suffix := time.Now().UnixNano()
	ownerID := insertAuthzUser(t, ctx, pool, suffix, "owner")
	reviewerA := insertAuthzUser(t, ctx, pool, suffix, "ra")
	reviewerB := insertAuthzUser(t, ctx, pool, suffix, "rb")

	var projectID int64
	if err := pool.QueryRow(ctx, `
		INSERT INTO projects (code, name, status, creator_user_id)
		VALUES ($1, $2, 'active', $3)
		RETURNING id
	`, fmt.Sprintf("duty-race-%d", suffix), "concurrent reviewer", ownerID).Scan(&projectID); err != nil {
		t.Fatalf("insert project: %v", err)
	}
	t.Cleanup(func() {
		_, _ = pool.Exec(context.Background(), `DELETE FROM project_members WHERE project_id = $1`, projectID)
		_, _ = pool.Exec(context.Background(), `DELETE FROM projects WHERE id = $1`, projectID)
		_, _ = pool.Exec(context.Background(), `DELETE FROM users WHERE id IN ($1,$2,$3)`, ownerID, reviewerA, reviewerB)
	})

	a := NewAuthorizer(pool)
	dutyAuthor := docduty.DutyAuthor
	if err := a.UpsertProjectMemberSpec(ctx, projectID, ownerID, MemberUpsert{
		ProjectRole: RoleOwner, DocumentDuty: &dutyAuthor,
	}); err != nil {
		t.Fatal(err)
	}
	if err := a.UpsertProjectMember(ctx, projectID, reviewerA, RoleMember); err != nil {
		t.Fatal(err)
	}
	if err := a.UpsertProjectMember(ctx, projectID, reviewerB, RoleMember); err != nil {
		t.Fatal(err)
	}

	duty := docduty.DutyReviewer
	var wg sync.WaitGroup
	errCh := make(chan error, 16)
	targets := []int64{reviewerA, reviewerB}
	for i := 0; i < 12; i++ {
		target := targets[i%2]
		wg.Add(1)
		go func(uid int64) {
			defer wg.Done()
			if err := a.UpsertProjectMemberSpec(ctx, projectID, uid, MemberUpsert{
				ProjectRole: RoleMember, DocumentDuty: &duty,
			}); err != nil {
				errCh <- err
			}
		}(target)
	}
	wg.Wait()
	close(errCh)
	for upsertErr := range errCh {
		t.Fatalf("concurrent upsert failed after retry: %v", upsertErr)
	}

	rows, err := pool.Query(ctx, `
		SELECT user_id FROM project_members
		WHERE project_id = $1 AND status = 'active' AND document_duty = 'reviewer'
		ORDER BY user_id
	`, projectID)
	if err != nil {
		t.Fatal(err)
	}
	defer rows.Close()
	var holders []int64
	for rows.Next() {
		var id int64
		if err := rows.Scan(&id); err != nil {
			t.Fatal(err)
		}
		holders = append(holders, id)
	}
	if err := rows.Err(); err != nil {
		t.Fatal(err)
	}
	if len(holders) != 1 {
		t.Fatalf("want exactly one active reviewer, got %v", holders)
	}
	if holders[0] != reviewerA && holders[0] != reviewerB {
		t.Fatalf("unexpected reviewer %d", holders[0])
	}
}

func TestUpsertProjectMember_AtomicDutyRevoke_PG(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping PostgreSQL duty revoke atomicity in -short")
	}
	dsn := db.SkipIfNoTestDatabaseURL(t)
	ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
	defer cancel()
	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		t.Fatalf("postgres unavailable: %v", err)
	}
	defer pool.Close()
	if err := pool.Ping(ctx); err != nil {
		t.Fatalf("postgres ping failed: %v", err)
	}
	if err := db.RunMigrationsOnceForTest(ctx, pool, dsn); err != nil {
		t.Fatalf("migration failed: %v", err)
	}

	suffix := time.Now().UnixNano()
	authorID := insertAuthzUser(t, ctx, pool, suffix, "author")
	reviewerA := insertAuthzUser(t, ctx, pool, suffix, "ra")
	reviewerB := insertAuthzUser(t, ctx, pool, suffix, "rb")

	var projectID int64
	if err := pool.QueryRow(ctx, `
		INSERT INTO projects (code, name, status, creator_user_id)
		VALUES ($1, $2, 'active', $3)
		RETURNING id
	`, fmt.Sprintf("duty-atomic-%d", suffix), "atomic upsert revoke", authorID).Scan(&projectID); err != nil {
		t.Fatalf("insert project: %v", err)
	}
	t.Cleanup(func() {
		_, _ = pool.Exec(context.Background(), `DELETE FROM document_records WHERE project_id = $1`, projectID)
		_, _ = pool.Exec(context.Background(), `DELETE FROM project_members WHERE project_id = $1`, projectID)
		_, _ = pool.Exec(context.Background(), `DELETE FROM projects WHERE id = $1`, projectID)
		_, _ = pool.Exec(context.Background(), `DELETE FROM users WHERE id IN ($1,$2,$3)`, authorID, reviewerA, reviewerB)
	})

	a := NewAuthorizer(pool)
	dutyAuthor := docduty.DutyAuthor
	dutyReviewer := docduty.DutyReviewer
	dutyNone := docduty.DutyNone
	if err := a.UpsertProjectMemberSpec(ctx, projectID, authorID, MemberUpsert{
		ProjectRole: RoleOwner, DocumentDuty: &dutyAuthor,
	}); err != nil {
		t.Fatal(err)
	}
	if err := a.UpsertProjectMemberSpec(ctx, projectID, reviewerA, MemberUpsert{
		ProjectRole: RoleMember, DocumentDuty: &dutyReviewer,
	}); err != nil {
		t.Fatal(err)
	}
	if err := a.UpsertProjectMemberSpec(ctx, projectID, reviewerB, MemberUpsert{
		ProjectRole: RoleMember, DocumentDuty: &dutyNone,
	}); err != nil {
		t.Fatal(err)
	}

	docID := fmt.Sprintf("doc_proj%d_intake", projectID)
	if _, err := pool.Exec(ctx, `
		INSERT INTO document_records (
			doc_id, project_id, task_key, title, stage, version,
			current_state, is_locked, is_template_only,
			author_id, author_name, reviewer_id, reviewer_name, approver_id, approver_name,
			created_at, updated_at
		) VALUES (
			$1, $2, 'intake', 'intake doc', 'V0.5', 1,
			'drafting', false, false,
			$3, '编制工程师', $4, '审核负责人', NULL, '',
			now(), now()
		)
	`, docID, projectID, authorID, reviewerA); err != nil {
		t.Fatalf("insert doc: %v", err)
	}

	if err := a.UpsertProjectMemberSpec(ctx, projectID, reviewerA, MemberUpsert{
		ProjectRole: RoleMember, DocumentDuty: &dutyNone,
	}); err != nil {
		t.Fatal(err)
	}
	assertAuthzDocSlot(t, ctx, pool, docID, "reviewer_id", 0)
	assertAuthzDocSlot(t, ctx, pool, docID, "author_id", authorID)
	assertAuthzMemberDuty(t, ctx, pool, projectID, reviewerA, docduty.DutyNone)

	if err := a.UpsertProjectMemberSpec(ctx, projectID, reviewerA, MemberUpsert{
		ProjectRole: RoleMember, DocumentDuty: &dutyReviewer,
	}); err != nil {
		t.Fatal(err)
	}
	if _, _, err := docduty.HealProjectDocuments(ctx, pool, projectID); err != nil {
		t.Fatal(err)
	}
	assertAuthzDocSlot(t, ctx, pool, docID, "reviewer_id", reviewerA)

	if err := a.UpsertProjectMemberSpec(ctx, projectID, reviewerB, MemberUpsert{
		ProjectRole: RoleMember, DocumentDuty: &dutyReviewer,
	}); err != nil {
		t.Fatal(err)
	}
	assertAuthzDocSlot(t, ctx, pool, docID, "reviewer_id", 0)
	assertAuthzDocSlot(t, ctx, pool, docID, "author_id", authorID)
	assertAuthzMemberDuty(t, ctx, pool, projectID, reviewerA, docduty.DutyNone)
	assertAuthzMemberDuty(t, ctx, pool, projectID, reviewerB, docduty.DutyReviewer)

	if _, _, err := docduty.HealProjectDocuments(ctx, pool, projectID); err != nil {
		t.Fatal(err)
	}
	assertAuthzDocSlot(t, ctx, pool, docID, "reviewer_id", reviewerB)
	assertAuthzDocSlot(t, ctx, pool, docID, "author_id", authorID)

	if err := a.UpsertProjectMemberSpec(ctx, projectID, reviewerB, MemberUpsert{
		ProjectRole: RoleMember, DocumentDuty: &dutyReviewer,
	}); err != nil {
		t.Fatal(err)
	}
	assertAuthzDocSlot(t, ctx, pool, docID, "reviewer_id", reviewerB)

	if err := a.SetProjectMemberStatus(ctx, projectID, reviewerB, StatusExited); err != nil {
		t.Fatal(err)
	}
	assertAuthzDocSlot(t, ctx, pool, docID, "reviewer_id", 0)
	assertAuthzDocSlot(t, ctx, pool, docID, "author_id", authorID)

	if err := a.SetProjectMemberStatus(ctx, projectID, reviewerB, StatusExited); err != nil {
		t.Fatal(err)
	}
	assertAuthzDocSlot(t, ctx, pool, docID, "reviewer_id", 0)
	assertAuthzDocSlot(t, ctx, pool, docID, "author_id", authorID)
}

func assertAuthzDocSlot(t *testing.T, ctx context.Context, pool *pgxpool.Pool, docID, column string, want int64) {
	t.Helper()
	var got int64
	q := fmt.Sprintf(`SELECT COALESCE(%s, 0) FROM document_records WHERE doc_id = $1`, column)
	if err := pool.QueryRow(ctx, q, docID).Scan(&got); err != nil {
		t.Fatal(err)
	}
	if got != want {
		t.Fatalf("%s %s=%d want %d", docID, column, got, want)
	}
}

func assertAuthzMemberDuty(t *testing.T, ctx context.Context, pool *pgxpool.Pool, projectID, userID int64, want string) {
	t.Helper()
	var got string
	if err := pool.QueryRow(ctx, `
		SELECT COALESCE(document_duty, '') FROM project_members WHERE project_id = $1 AND user_id = $2
	`, projectID, userID).Scan(&got); err != nil {
		t.Fatal(err)
	}
	if got != want {
		t.Fatalf("user %d document_duty=%q want %q", userID, got, want)
	}
}

func insertAuthzUser(t *testing.T, ctx context.Context, pool *pgxpool.Pool, suffix int64, kind string) int64 {
	t.Helper()
	var id int64
	if err := pool.QueryRow(ctx, `
		INSERT INTO users (account_id, name, role)
		VALUES ($1, $2, 'engineer')
		RETURNING id
	`, fmt.Sprintf("authz_%s_%d", kind, suffix), "用户"+kind).Scan(&id); err != nil {
		t.Fatalf("insert user %s: %v", kind, err)
	}
	return id
}