RenewCMS/test/accept_invitation_atomic_test.go

package test

import (
	"RenewCMS/api/handlers/auth"
	validator "RenewCMS/api/pkg/validator"
	"RenewCMS/internal/domain/invite"
	"RenewCMS/internal/infrastructure/persistence/models"
	persistence "RenewCMS/internal/infrastructure/persistence/repositories"
	"RenewCMS/internal/infrastructure/useCases"
	"bytes"
	"fmt"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/glebarez/sqlite"
	"github.com/go-chi/chi/v5"
	"github.com/go-chi/jwtauth/v5"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

const atomicAcceptRuns = 50

type atomicPlainHasher struct{}

func (atomicPlainHasher) Hash(value string) (string, error) {
	return "hashed:" + value, nil
}

func (atomicPlainHasher) Compare(hashed string, value string) error {
	if hashed == "hashed:"+value {
		return nil
	}
	return fmt.Errorf("hash mismatch")
}

type atomicInviteRepoDecorator struct {
	invite.Repository
	afterGetAll func()
}

func (d *atomicInviteRepoDecorator) GetAll() []invite.Invite {
	all := d.Repository.GetAll()
	d.afterGetAll()
	return all
}

type atomicLookupBarrier struct {
	mu      sync.Mutex
	arrived int
	parties int
	open    chan struct{}
}

func newAtomicLookupBarrier(parties int) *atomicLookupBarrier {
	return &atomicLookupBarrier{parties: parties, open: make(chan struct{})}
}

func (b *atomicLookupBarrier) wait() {
	b.mu.Lock()
	b.arrived++
	if b.arrived == b.parties {
		close(b.open)
	}
	b.mu.Unlock()
	select {
	case <-b.open:
	case <-time.After(5 * time.Second):
	}
}

func withLookupBarrier(parties int) func(db *gorm.DB, repo invite.Repository) invite.Repository {
	barrier := newAtomicLookupBarrier(parties)
	return func(db *gorm.DB, repo invite.Repository) invite.Repository {
		return &atomicInviteRepoDecorator{Repository: repo, afterGetAll: barrier.wait}
	}
}

type atomicAcceptFixture struct {
	db     *gorm.DB
	router *chi.Mux
}

func openRealSqlite(t *testing.T) *gorm.DB {
	t.Helper()
	path := filepath.ToSlash(filepath.Join(t.TempDir(), "atomic.db"))
	dsn := "file:" + path + "?_pragma=busy_timeout(10000)&_pragma=journal_mode(WAL)"
	db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{Logger: logger.Discard})
	if err != nil {
		t.Fatalf("open sqlite: %v", err)
	}
	if err := db.AutoMigrate(&models.Article{}, &models.User{}, &models.Invite{}); err != nil {
		t.Fatalf("migrate: %v", err)
	}
	t.Cleanup(func() {
		if sqlDB, err := db.DB(); err == nil {
			_ = sqlDB.Close()
		}
	})
	return db
}

func newAtomicAcceptFixture(t *testing.T, wrap func(db *gorm.DB, repo invite.Repository) invite.Repository) *atomicAcceptFixture {
	t.Helper()
	validator.InitValidator()
	db := openRealSqlite(t)
	users := persistence.NewUserRepository(db)
	var invites invite.Repository = persistence.NewInviteRepository(db)
	if wrap != nil {
		invites = wrap(db, invites)
	}
	hasher := atomicPlainHasher{}
	handler := auth.NewAcceptInviteHandler(
		useCases.NewGetInviteUseCase(invites, hasher),
		useCases.NewAcceptInviteUseCase(users, invites, hasher, discardRecorder),
		jwtauth.New("HS256", []byte("test-secret"), nil),
		true,
	)
	router := chi.NewRouter()
	router.Post("/v1/auth/invite/accept", handler.AcceptInvite)

	seedRealUser(t, db, models.User{Username: "marta", Password: "x", Email: "marta@renew.example", IsVerified: true})
	seedRealInvite(t, db)
	return &atomicAcceptFixture{db: db, router: router}
}

func seedRealUser(t *testing.T, db *gorm.DB, u models.User) models.User {
	t.Helper()
	if err := db.Create(&u).Error; err != nil {
		t.Fatalf("seed user: %v", err)
	}
	return u
}

func seedRealInvite(t *testing.T, db *gorm.DB) models.Invite {
	t.Helper()
	pending := models.Invite{
		Email:      acceptEmail,
		HashedCode: "hashed:" + acceptCode,
		ExpiresAt:  time.Now().Add(48 * time.Hour),
		InvitedBy:  1,
	}
	if err := db.Create(&pending).Error; err != nil {
		t.Fatalf("seed invite: %v", err)
	}
	return pending
}

func (f *atomicAcceptFixture) accept(username string) *httptest.ResponseRecorder {
	req := httptest.NewRequest(http.MethodPost, "/v1/auth/invite/accept", bytes.NewBufferString(acceptBody(acceptCode, username, acceptPassword)))
	rec := httptest.NewRecorder()
	f.router.ServeHTTP(rec, req)
	return rec
}

func (f *atomicAcceptFixture) usersWithEmail(email string) int64 {
	var count int64
	f.db.Model(&models.User{}).Where("email = ?", email).Count(&count)
	return count
}

func (f *atomicAcceptFixture) usersNamed(username string) int64 {
	var count int64
	f.db.Model(&models.User{}).Where("username = ?", username).Count(&count)
	return count
}

func (f *atomicAcceptFixture) storedInvite() models.Invite {
	var stored models.Invite
	f.db.First(&stored)
	return stored
}

func (f *atomicAcceptFixture) acceptConcurrently(usernames ...string) []*httptest.ResponseRecorder {
	recorders := make([]*httptest.ResponseRecorder, len(usernames))
	start := make(chan struct{})
	var wg sync.WaitGroup
	for index, username := range usernames {
		wg.Add(1)
		go func() {
			defer wg.Done()
			<-start
			recorders[index] = f.accept(username)
		}()
	}
	close(start)
	wg.Wait()
	return recorders
}

func assertOneWinnerOneLoser(t *testing.T, f *atomicAcceptFixture, recorders []*httptest.ResponseRecorder) {
	t.Helper()
	winners, losers := 0, 0
	for _, rec := range recorders {
		switch {
		case rec.Code == http.StatusOK && jwtCookieOf(rec) != nil:
			winners++
		case rec.Code == http.StatusNotFound && rec.Body.String() == inviteInvalidBody && jwtCookieOf(rec) == nil:
			losers++
		default:
			t.Errorf("unexpected response %d %q cookie=%v", rec.Code, rec.Body.String(), jwtCookieOf(rec))
		}
	}
	if winners != 1 || losers != 1 {
		t.Errorf("expected one winner and one loser, got %d and %d", winners, losers)
	}
	if count := f.usersWithEmail(acceptEmail); count != 1 {
		t.Errorf("expected exactly one user for the invited email, got %d", count)
	}
	if f.storedInvite().AcceptedAt == nil {
		t.Errorf("invite must have accepted_at set")
	}
}

func TestAcceptInvitationAtomic(t *testing.T) {
	t.Run("Two concurrent accepts of one code create exactly one user", func(t *testing.T) {
		f := newAtomicAcceptFixture(t, withLookupBarrier(2))

		assertOneWinnerOneLoser(t, f, f.acceptConcurrently("lea", "lea2"))
	})

	t.Run("The concurrency test is repeatable", func(t *testing.T) {
		for run := range atomicAcceptRuns {
			t.Run(fmt.Sprintf("run %d", run), func(t *testing.T) {
				f := newAtomicAcceptFixture(t, withLookupBarrier(2))

				assertOneWinnerOneLoser(t, f, f.acceptConcurrently("lea", "lea2"))
			})
		}
	})

	t.Run("A used code loses at the conditional update, not at the earlier lookup", func(t *testing.T) {
		f := newAtomicAcceptFixture(t, func(db *gorm.DB, repo invite.Repository) invite.Repository {
			return &atomicInviteRepoDecorator{Repository: repo, afterGetAll: func() {
				db.Exec("UPDATE invites SET accepted_at = ?", time.Now())
			}}
		})

		rec := f.accept("lea")

		if rec.Code != http.StatusNotFound || rec.Body.String() != inviteInvalidBody {
			t.Errorf("expected 404 %q, got %d %q", inviteInvalidBody, rec.Code, rec.Body.String())
		}
		if jwtCookieOf(rec) != nil {
			t.Errorf("no jwt cookie may be set")
		}
		if f.usersNamed("lea") != 0 {
			t.Errorf("no user may be inserted for the losing accept")
		}
	})

	t.Run("A failed user insert leaves the invite pending", func(t *testing.T) {
		f := newAtomicAcceptFixture(t, nil)

		rec := f.accept("marta")

		if rec.Code != http.StatusConflict || rec.Body.String() != usernameTakenBody {
			t.Errorf("expected 409 %q, got %d %q", usernameTakenBody, rec.Code, rec.Body.String())
		}
		if f.storedInvite().AcceptedAt != nil {
			t.Errorf("invite must stay pending")
		}
	})

	t.Run("A user insert rejected by the database leaves the invite pending", func(t *testing.T) {
		f := newAtomicAcceptFixture(t, nil)
		seedRealUser(t, f.db, models.User{Username: "leam", Password: "x", Email: "someone@renew.example"})
		if err := f.db.Exec("CREATE TRIGGER users_reject_insert BEFORE INSERT ON users WHEN NEW.username = 'lea' BEGIN SELECT RAISE(ABORT, 'insert refused'); END").Error; err != nil {
			t.Fatalf("trigger: %v", err)
		}

		rec := f.accept("lea")

		if rec.Code != http.StatusInternalServerError || jwtCookieOf(rec) != nil {
			t.Errorf("expected 500 without cookie, got %d", rec.Code)
		}
		if f.storedInvite().AcceptedAt != nil {
			t.Errorf("invite must stay pending after the insert failed")
		}
		if f.usersNamed("lea") != 0 {
			t.Errorf("no user lea may exist")
		}
	})

	t.Run("A failed invite update leaves no user behind", func(t *testing.T) {
		f := newAtomicAcceptFixture(t, nil)
		if err := f.db.Exec("CREATE TRIGGER invites_frozen BEFORE UPDATE ON invites BEGIN SELECT RAISE(ABORT, 'invites are frozen'); END").Error; err != nil {
			t.Fatalf("trigger: %v", err)
		}

		rec := f.accept("lea")

		if rec.Code != http.StatusInternalServerError || jwtCookieOf(rec) != nil {
			t.Errorf("expected 500 without cookie, got %d", rec.Code)
		}
		if f.usersNamed("lea") != 0 {
			t.Errorf("no user lea may exist")
		}
		if f.storedInvite().AcceptedAt != nil {
			t.Errorf("invite must stay pending")
		}
	})
}