proxy-pool/internal/domain/workerruntime/memory_test.go

package workerruntime

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

func TestMemoryStoreReplacesSparseRuntimeAndClearsMissingCounters(t *testing.T) {
	now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	ctx := context.Background()
	registerRuntimeSession(t, store, "session-a", time.Minute)
	report := Report{
		WorkerID: "worker-a", SessionID: "session-a", Sequence: 1,
		SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now,
		Counters: []Counter{{ProxyID: "proxy-a", Active: 2, Reserved: 1}},
	}
	if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(first): %v", err)
	}
	got, err := store.ReadRuntime(ctx, []OwnedProxy{{ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 9}})
	if err != nil || len(got) != 1 || got[0] != (Snapshot{ProxyID: "proxy-a", Active: 2, Reserved: 1, Fresh: true}) {
		t.Fatalf("ReadRuntime(first) = %+v, %v", got, err)
	}

	report.Sequence = 2
	report.ObservedAt = now.Add(time.Second)
	report.Counters = nil
	if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(empty): %v", err)
	}
	got, err = store.ReadRuntime(ctx, []OwnedProxy{{ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 9}})
	if err != nil || len(got) != 1 || got[0] != (Snapshot{ProxyID: "proxy-a", Fresh: true}) {
		t.Fatalf("ReadRuntime(empty) = %+v, %v", got, err)
	}
}

func TestMemoryStoreFencesSessionsAndReportSequence(t *testing.T) {
	now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	ctx := context.Background()
	registerRuntimeSession(t, store, "session-a", time.Minute)
	report := Report{
		WorkerID: "worker-a", SessionID: "session-a", Sequence: 2,
		SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now,
		Counters: []Counter{{ProxyID: "proxy-a", Active: 1}},
	}
	if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(first): %v", err)
	}
	if err := store.ReplaceRuntime(ctx, report, time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(idempotent): %v", err)
	}
	conflict := report
	conflict.Counters = []Counter{{ProxyID: "proxy-a", Active: 2}}
	if err := store.ReplaceRuntime(ctx, conflict, time.Minute); !errors.Is(err, ErrConflictingReport) {
		t.Fatalf("ReplaceRuntime(conflict) error = %v", err)
	}
	stale := report
	stale.Sequence = 1
	if err := store.ReplaceRuntime(ctx, stale, time.Minute); !errors.Is(err, ErrStaleReport) {
		t.Fatalf("ReplaceRuntime(stale) error = %v", err)
	}
	registerRuntimeSession(t, store, "session-b", time.Minute)
	if err := store.ReplaceRuntime(ctx, report, time.Minute); !errors.Is(err, ErrStaleSession) {
		t.Fatalf("ReplaceRuntime(old session) error = %v", err)
	}
}

func TestMemoryStoreFailsClosedForExpiredOrOlderOwnershipReport(t *testing.T) {
	now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	ctx := context.Background()
	registerRuntimeSession(t, store, "session-a", time.Minute)
	if err := store.ReplaceRuntime(ctx, Report{
		WorkerID: "worker-a", SessionID: "session-a", Sequence: 1,
		SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now,
	}, time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(): %v", err)
	}
	queries := []OwnedProxy{
		{ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 10},
		{ProxyID: "proxy-b", WorkerID: "worker-a", OwnershipEpoch: 9},
	}
	got, err := store.ReadRuntime(ctx, queries)
	if err != nil || got[0].Fresh || !got[1].Fresh {
		t.Fatalf("ReadRuntime(ownership fence) = %+v, %v", got, err)
	}
	now = now.Add(time.Minute)
	got, err = store.ReadRuntime(ctx, queries[1:])
	if err != nil || len(got) != 1 || got[0].Fresh {
		t.Fatalf("ReadRuntime(expired) = %+v, %v", got, err)
	}
}

func TestMemoryStoreRejectsRuntimeBeyondAcknowledgedSnapshot(t *testing.T) {
	now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	registerRuntimeSession(t, store, "session-a", time.Minute)
	report := Report{
		WorkerID: "worker-a", SessionID: "session-a", Sequence: 1,
		SnapshotVersion: 4, OwnershipEpoch: 9, ObservedAt: now,
	}
	if err := store.ReplaceRuntime(context.Background(), report, time.Minute); !errors.Is(err, ErrStaleReport) {
		t.Fatalf("ReplaceRuntime(ahead snapshot) error = %v, want ErrStaleReport", err)
	}
	report.SnapshotVersion = 3
	report.OwnershipEpoch = 10
	if err := store.ReplaceRuntime(context.Background(), report, time.Minute); !errors.Is(err, ErrStaleReport) {
		t.Fatalf("ReplaceRuntime(ahead epoch) error = %v, want ErrStaleReport", err)
	}
}

func TestMemoryStoreExpiredSameIdentitySessionDoesNotReactivateOldReport(t *testing.T) {
	now := time.Date(2026, 7, 30, 15, 0, 0, 0, time.UTC)
	store := newRuntimeStore(t, &now)
	registerRuntimeSession(t, store, "session-a", time.Second)
	if err := store.ReplaceRuntime(context.Background(), Report{
		WorkerID: "worker-a", SessionID: "session-a", Sequence: 1,
		SnapshotVersion: 3, OwnershipEpoch: 9, ObservedAt: now,
	}, time.Minute); err != nil {
		t.Fatalf("ReplaceRuntime(): %v", err)
	}
	session := store.sessions["worker-a"]
	session.expiresAt = now.Add(time.Second)
	store.sessions["worker-a"] = session
	now = now.Add(2 * time.Second)
	registerRuntimeSession(t, store, "session-a", time.Minute)
	got, err := store.ReadRuntime(context.Background(), []OwnedProxy{{
		ProxyID: "proxy-a", WorkerID: "worker-a", OwnershipEpoch: 9,
	}})
	if err != nil || len(got) != 1 || got[0].Fresh {
		t.Fatalf("ReadRuntime(after re-register) = %+v, %v; want stale", got, err)
	}
}

func newRuntimeStore(t *testing.T, now *time.Time) *MemoryStore {
	t.Helper()
	store, err := NewMemoryStore(func() time.Time { return *now })
	if err != nil {
		t.Fatalf("NewMemoryStore(): %v", err)
	}
	return store
}

func registerRuntimeSession(t *testing.T, store *MemoryStore, sessionID string, ttl time.Duration) {
	t.Helper()
	if err := store.ReplaceSession(context.Background(), Session{
		WorkerID: "worker-a", InstanceID: "instance-a", SessionID: sessionID,
		AckedSnapshotVersion: 3, AckedOwnershipEpoch: 9,
	}, ttl); err != nil {
		t.Fatalf("ReplaceSession(): %v", err)
	}
}