easyai-ai-gateway/apps/api/internal/securityevents/retirement_worker_test.go

package securityevents

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
	"github.com/google/uuid"
)

func TestSecurityEventRetirementWorkerOutlivesCanceledRuntime(t *testing.T) {
	connection := retiringWorkerTestConnection(time.Now().UTC())
	repository := &memoryConnectionRepository{connection: &connection}
	secrets := &memorySecretStore{values: retirementWorkerTestSecrets(connection)}

	runtimeCtx, cancelRuntime := context.WithCancel(context.Background())
	manager := &ConnectionManager{
		ctx: runtimeCtx, repository: repository, secrets: secrets,
		config: ConnectionManagerConfig{RetirementDuration: 40 * time.Millisecond},
	}
	go manager.finishRetirement(connection)
	cancelRuntime()
	time.Sleep(10 * time.Millisecond)
	if _, err := repository.SecurityEventConnection(context.Background()); err != nil {
		t.Fatalf("runtime cancellation unexpectedly finalized retirement: %v", err)
	}

	serverCtx, cancelServer := context.WithCancel(context.Background())
	defer cancelServer()
	go runSecurityEventRetirementWorker(serverCtx, repository, 40*time.Millisecond, 5*time.Millisecond)
	waitForRetirementFinalization(t, repository)

	for reference := range retirementWorkerTestSecrets(connection) {
		repository.mutex.Lock()
		_, queued := repository.cleanupDue[reference]
		repository.mutex.Unlock()
		if !queued {
			t.Fatalf("retired Secret %q was not durably queued", reference)
		}
		if !secrets.Has(reference) {
			t.Fatalf("retirement worker directly deleted Secret %q before the durable cleanup worker claimed it", reference)
		}
	}
}

func TestSecurityEventRetirementWorkerResumesAfterServerRestart(t *testing.T) {
	connection := retiringWorkerTestConnection(time.Now().UTC())
	repository := &memoryConnectionRepository{connection: &connection}

	firstCtx, stopFirst := context.WithCancel(context.Background())
	go runSecurityEventRetirementWorker(firstCtx, repository, 80*time.Millisecond, 5*time.Millisecond)
	time.Sleep(15 * time.Millisecond)
	stopFirst()
	if _, err := repository.SecurityEventConnection(context.Background()); err != nil {
		t.Fatalf("not-yet-due retirement disappeared before restart: %v", err)
	}

	time.Sleep(75 * time.Millisecond)
	secondCtx, stopSecond := context.WithCancel(context.Background())
	defer stopSecond()
	go runSecurityEventRetirementWorker(secondCtx, repository, 80*time.Millisecond, 5*time.Millisecond)
	waitForRetirementFinalization(t, repository)
}

func TestConcurrentSecurityEventRetirementWorkersFinalizeOneGenerationOnce(t *testing.T) {
	connection := retiringWorkerTestConnection(time.Now().UTC().Add(-time.Minute))
	repository := &memoryConnectionRepository{connection: &connection}

	start := make(chan struct{})
	done := make(chan struct{}, 2)
	for range 2 {
		go func() {
			<-start
			finalizeDueSecurityEventRetirement(context.Background(), repository, time.Millisecond)
			done <- struct{}{}
		}()
	}
	close(start)
	<-done
	<-done

	if _, err := repository.SecurityEventConnection(context.Background()); !errors.Is(err, store.ErrSecurityEventConnectionNotFound) {
		t.Fatalf("current retiring generation still exists: %v", err)
	}
	repository.mutex.Lock()
	finalizeCalls := repository.finalizeCalls
	queued := len(repository.cleanupDue)
	repository.mutex.Unlock()
	if finalizeCalls != 1 {
		t.Fatalf("retirement generation finalized %d times, want 1", finalizeCalls)
	}
	if queued != len(retirementWorkerTestSecrets(connection)) {
		t.Fatalf("queued Secret references=%d want=%d", queued, len(retirementWorkerTestSecrets(connection)))
	}
}

func TestSecurityEventRetirementWorkerPreservesChangedGeneration(t *testing.T) {
	connection := retiringWorkerTestConnection(time.Now().UTC().Add(-time.Minute))
	initialVersion := connection.Version
	repository := &memoryConnectionRepository{connection: &connection}
	repository.beforeFinalize = func(current *store.SecurityEventConnection) {
		current.LifecycleStatus = "enabled"
		current.Version++
	}

	finalizeDueSecurityEventRetirement(context.Background(), repository, time.Millisecond)

	current, err := repository.SecurityEventConnection(context.Background())
	if err != nil || current.LifecycleStatus != "enabled" || current.Version != initialVersion+1 {
		t.Fatalf("changed generation was not preserved: lifecycle=%q version=%d err=%v", current.LifecycleStatus, current.Version, err)
	}
	repository.mutex.Lock()
	queued := len(repository.cleanupDue)
	repository.mutex.Unlock()
	if queued != 0 {
		t.Fatalf("changed generation queued %d active Secret references", queued)
	}
}

func TestSecurityEventRetirementWorkerKeepsConnectionWhenCleanupQueueCannotAdoptReferences(t *testing.T) {
	connection := retiringWorkerTestConnection(time.Now().UTC().Add(-time.Minute))
	repository := &memoryConnectionRepository{
		connection:    &connection,
		cleanupClaims: map[string]string{connection.CredentialRef: uuid.NewString()},
	}

	finalizeDueSecurityEventRetirement(context.Background(), repository, time.Millisecond)

	current, err := repository.SecurityEventConnection(context.Background())
	if err != nil || current.ConnectionID != connection.ConnectionID || current.LifecycleStatus != "retiring" {
		t.Fatalf("failed durable handoff lost retiring connection: lifecycle=%q err=%v", current.LifecycleStatus, err)
	}
	repository.mutex.Lock()
	finalizeCalls := repository.finalizeCalls
	repository.mutex.Unlock()
	if finalizeCalls != 0 {
		t.Fatalf("failed durable handoff finalized generation %d times", finalizeCalls)
	}
}

func retiringWorkerTestConnection(updatedAt time.Time) store.SecurityEventConnection {
	managementReference := "ssf-management-retirement-worker-" + uuid.NewString()
	nextReference := "ssf-push-next-retirement-worker-" + uuid.NewString()
	return store.SecurityEventConnection{
		ConnectionID:            uuid.NewString(),
		CredentialRef:           "ssf-push-retirement-worker-" + uuid.NewString(),
		NextCredentialRef:       &nextReference,
		ManagementCredentialRef: &managementReference,
		LifecycleStatus:         "retiring",
		Version:                 9,
		UpdatedAt:               updatedAt,
	}
}

func retirementWorkerTestSecrets(connection store.SecurityEventConnection) map[string][]byte {
	secrets := map[string][]byte{connection.CredentialRef: []byte("retired-push-secret-value")}
	if connection.NextCredentialRef != nil {
		secrets[*connection.NextCredentialRef] = []byte("retired-next-push-secret-value")
	}
	if connection.ManagementCredentialRef != nil {
		secrets[*connection.ManagementCredentialRef] = []byte("retired-management-secret-value")
	}
	return secrets
}

func waitForRetirementFinalization(t *testing.T, repository *memoryConnectionRepository) {
	t.Helper()
	deadline := time.Now().Add(time.Second)
	for time.Now().Before(deadline) {
		if _, err := repository.SecurityEventConnection(context.Background()); errors.Is(err, store.ErrSecurityEventConnectionNotFound) {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatal("retiring connection was not finalized by the server-lifecycle worker")
}