proxy-pool/internal/domain/workerruntime/contracttest/contract.go

package contracttest

import (
	"context"
	"crypto/sha256"
	"errors"
	"testing"
	"time"

	"proxy-pool/internal/domain/workerruntime"
)

type Fixture struct {
	Store   workerruntime.ControlStore
	Reader  workerruntime.RuntimeReader
	TTL     time.Duration
	Advance func(time.Duration)
}

type Factory func(*testing.T) Fixture

// Run exercises the public control-store behavior shared by Memory and Redis.
func Run(t *testing.T, factory Factory) {
	t.Helper()
	t.Run("acknowledged runtime lifecycle", func(t *testing.T) { runLifecycle(t, newFixture(t, factory)) })
	t.Run("negative acknowledgement fences runtime", func(t *testing.T) { runNegativeAck(t, newFixture(t, factory)) })
}

func runLifecycle(t *testing.T, fixture Fixture) {
	t.Helper()
	ctx := context.Background()
	open(t, fixture.Store, fixture.TTL)
	epoch, err := fixture.Store.CurrentOwnershipEpoch(ctx)
	if err != nil {
		t.Fatalf("CurrentOwnershipEpoch(): %v", err)
	}
	reference := snapshot(7, epoch, "snapshot-7")
	report := runtimeReport(1, 7, epoch)
	if err := fixture.Store.ReplaceRuntime(ctx, report, fixture.TTL); !errors.Is(err, workerruntime.ErrSnapshotMismatch) {
		t.Fatalf("ReplaceRuntime(before ACK) error = %v", err)
	}
	if err := fixture.Store.RecordIssuedSnapshot(ctx, reference, fixture.TTL); err != nil {
		t.Fatalf("RecordIssuedSnapshot(): %v", err)
	}
	ack := workerruntime.SnapshotAcknowledgement{WorkerID: "worker-a", SessionID: "session-a", Reference: reference, Applied: true}
	if err := fixture.Store.AcknowledgeSnapshot(ctx, ack, fixture.TTL); err != nil {
		t.Fatalf("AcknowledgeSnapshot(): %v", err)
	}
	if err := fixture.Store.ReplaceRuntime(ctx, report, fixture.TTL); err != nil {
		t.Fatalf("ReplaceRuntime(after ACK): %v", err)
	}
	assertFresh(t, fixture.Reader, epoch, true)
	if err := fixture.Store.AcknowledgeSnapshot(ctx, ack, fixture.TTL); err != nil {
		t.Fatalf("AcknowledgeSnapshot(replay): %v", err)
	}
	if err := fixture.Store.ReplaceRuntime(ctx, runtimeReport(0, 7, epoch), fixture.TTL); !errors.Is(err, workerruntime.ErrInvalidReport) {
		t.Fatalf("ReplaceRuntime(invalid sequence) error = %v", err)
	}
	fixture.Advance(2 * fixture.TTL)
	assertFresh(t, fixture.Reader, epoch, false)
}

func runNegativeAck(t *testing.T, fixture Fixture) {
	t.Helper()
	ctx := context.Background()
	open(t, fixture.Store, fixture.TTL)
	epoch, err := fixture.Store.CurrentOwnershipEpoch(ctx)
	if err != nil {
		t.Fatalf("CurrentOwnershipEpoch(): %v", err)
	}
	first := snapshot(7, epoch, "snapshot-7")
	if err := fixture.Store.RecordIssuedSnapshot(ctx, first, fixture.TTL); err != nil {
		t.Fatalf("RecordIssuedSnapshot(first): %v", err)
	}
	if err := fixture.Store.AcknowledgeSnapshot(ctx, workerruntime.SnapshotAcknowledgement{WorkerID: "worker-a", SessionID: "session-a", Reference: first, Applied: true}, fixture.TTL); err != nil {
		t.Fatalf("AcknowledgeSnapshot(first): %v", err)
	}
	second := snapshot(8, epoch, "snapshot-8")
	if err := fixture.Store.RecordIssuedSnapshot(ctx, second, fixture.TTL); err != nil {
		t.Fatalf("RecordIssuedSnapshot(second): %v", err)
	}
	if err := fixture.Store.AcknowledgeSnapshot(ctx, workerruntime.SnapshotAcknowledgement{
		WorkerID: "worker-a", SessionID: "session-a", Reference: second, ErrorCode: "apply_failed",
	}, fixture.TTL); err != nil {
		t.Fatalf("AcknowledgeSnapshot(negative): %v", err)
	}
	if err := fixture.Store.ReplaceRuntime(ctx, runtimeReport(1, 7, epoch), fixture.TTL); !errors.Is(err, workerruntime.ErrSnapshotMismatch) {
		t.Fatalf("ReplaceRuntime(delayed): %v", err)
	}
}

func newFixture(t *testing.T, factory Factory) Fixture {
	t.Helper()
	fixture := factory(t)
	if fixture.Store == nil || fixture.Reader == nil || fixture.TTL <= 0 || fixture.Advance == nil {
		t.Fatal("contract fixture is incomplete")
	}
	return fixture
}

func open(t *testing.T, store workerruntime.ControlStore, ttl time.Duration) {
	t.Helper()
	err := store.OpenSession(context.Background(), workerruntime.Session{
		WorkerID: "worker-a", InstanceID: "instance-a", SessionID: "session-a", Zone: "zone-a", ProtocolVersion: 1,
	}, ttl)
	if err != nil {
		t.Fatalf("OpenSession(): %v", err)
	}
}

func snapshot(version, epoch uint64, value string) workerruntime.SnapshotReference {
	return workerruntime.SnapshotReference{
		WorkerID: "worker-a", Version: version, OwnershipEpoch: epoch, Checksum: sha256.Sum256([]byte(value)),
	}
}

func runtimeReport(sequence, version, epoch uint64) workerruntime.Report {
	return workerruntime.Report{
		WorkerID: "worker-a", SessionID: "session-a", Sequence: sequence, SnapshotVersion: version,
		OwnershipEpoch: epoch, ObservedAt: time.Date(2026, 7, 31, 9, 0, 0, 0, time.UTC),
	}
}

func assertFresh(t *testing.T, reader workerruntime.RuntimeReader, epoch uint64, want bool) {
	t.Helper()
	snapshots, err := reader.ReadRuntime(context.Background(), []workerruntime.OwnedProxy{{
		ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: epoch,
	}})
	if err != nil || len(snapshots) != 1 || snapshots[0].Fresh != want {
		t.Fatalf("ReadRuntime() = %+v, %v; want Fresh=%t", snapshots, err, want)
	}
}