Build AI gateway control plane and admin UI

1 files changed, 212 insertions, 0 deletions

diff --git a/internal/controlplane/manager_test.go b/internal/controlplane/manager_test.go
new file mode 100644
index 0000000..8dd4012
--- /dev/null
+++ b/internal/controlplane/manager_test.go
@@ -0,0 +1,212 @@
+package controlplane
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"log/slog"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"aigw/internal/auth"
+	"aigw/internal/catalog"
+)
+
+type fakeManagerStore struct {
+	redisEnabled       bool
+	snapshot           atomic.Value
+	databaseGeneration atomic.Int64
+	publishCalls       atomic.Int64
+	subscribeCalls     atomic.Int64
+	publishErr         error
+	published          chan ChangeEvent
+	subscribe          func(context.Context, int64) (<-chan ChangeMessage, func() error, error)
+}
+
+type safeLogBuffer struct {
+	mu  sync.Mutex
+	buf bytes.Buffer
+}
+
+func (b *safeLogBuffer) Write(data []byte) (int, error) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.buf.Write(data)
+}
+
+func (b *safeLogBuffer) String() string {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.buf.String()
+}
+
+func newFakeManagerStore(generation int64) *fakeManagerStore {
+	store := &fakeManagerStore{published: make(chan ChangeEvent, 8)}
+	store.snapshot.Store(Snapshot{Generation: generation})
+	store.databaseGeneration.Store(generation)
+	return store
+}
+
+func (s *fakeManagerStore) LoadSnapshot(context.Context) (Snapshot, error) {
+	return s.snapshot.Load().(Snapshot), nil
+}
+
+func (s *fakeManagerStore) DatabaseGeneration(context.Context) (int64, error) {
+	return s.databaseGeneration.Load(), nil
+}
+
+func (s *fakeManagerStore) PublishChange(_ context.Context, event ChangeEvent) error {
+	s.publishCalls.Add(1)
+	s.published <- event
+	return s.publishErr
+}
+
+func (s *fakeManagerStore) Subscribe(ctx context.Context) (<-chan ChangeMessage, func() error, error) {
+	call := s.subscribeCalls.Add(1)
+	if s.subscribe != nil {
+		return s.subscribe(ctx, call)
+	}
+	channel := make(chan ChangeMessage)
+	return channel, func() error { return nil }, nil
+}
+
+func (s *fakeManagerStore) RedisEnabled() bool {
+	return s.redisEnabled
+}
+
+func newTestManager(store managerStore, logger *slog.Logger, interval time.Duration) *Manager {
+	return NewManager(store, catalog.NewModels(nil), auth.NewDynamic(nil, false), logger, interval)
+}
+
+func TestAfterMutationReloadsLocallyWhenRedisPublishFails(t *testing.T) {
+	store := newFakeManagerStore(7)
+	store.redisEnabled = true
+	store.publishErr = errors.New("redis unavailable")
+	var logs safeLogBuffer
+	manager := newTestManager(store, slog.New(slog.NewTextHandler(&logs, nil)), 10*time.Millisecond)
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	go func() {
+		manager.Run(ctx)
+		close(done)
+	}()
+
+	if err := manager.AfterMutation(context.Background(), 7, "model", "model-1"); err != nil {
+		t.Fatalf("mutation unexpectedly failed: %v", err)
+	}
+	if manager.Generation() != 7 {
+		t.Fatalf("local generation = %d, want 7", manager.Generation())
+	}
+	select {
+	case event := <-store.published:
+		if event.Generation != 7 || event.Resource != "model" {
+			t.Fatalf("unexpected event: %+v", event)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("broadcast was not attempted")
+	}
+	waitUntil(t, time.Second, func() bool { return strings.Contains(logs.String(), "control_plane_publish_failed") })
+
+	cancel()
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("manager did not stop")
+	}
+}
+
+func TestPollingContinuesWhileRedisSubscribeIsBlocked(t *testing.T) {
+	store := newFakeManagerStore(1)
+	store.redisEnabled = true
+	store.subscribe = func(ctx context.Context, _ int64) (<-chan ChangeMessage, func() error, error) {
+		<-ctx.Done()
+		return nil, nil, ctx.Err()
+	}
+	manager := newTestManager(store, slog.New(slog.NewTextHandler(&safeLogBuffer{}, nil)), 10*time.Millisecond)
+	if _, err := manager.Reload(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	go func() {
+		manager.Run(ctx)
+		close(done)
+	}()
+	store.snapshot.Store(Snapshot{Generation: 2})
+	store.databaseGeneration.Store(2)
+	waitUntil(t, time.Second, func() bool { return manager.Generation() == 2 })
+
+	cancel()
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("manager did not stop")
+	}
+}
+
+func TestSubscriptionReconnectsAfterChannelCloses(t *testing.T) {
+	store := newFakeManagerStore(1)
+	store.redisEnabled = true
+	first := make(chan ChangeMessage)
+	second := make(chan ChangeMessage, 1)
+	store.subscribe = func(_ context.Context, call int64) (<-chan ChangeMessage, func() error, error) {
+		if call == 1 {
+			return first, func() error { return nil }, nil
+		}
+		return second, func() error { return nil }, nil
+	}
+	manager := newTestManager(store, slog.New(slog.NewTextHandler(&safeLogBuffer{}, nil)), 10*time.Millisecond)
+	if _, err := manager.Reload(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	go func() {
+		manager.Run(ctx)
+		close(done)
+	}()
+	waitUntil(t, time.Second, func() bool { return store.subscribeCalls.Load() == 1 })
+	close(first)
+	waitUntil(t, time.Second, func() bool { return store.subscribeCalls.Load() >= 2 && manager.RedisConnected() })
+	store.snapshot.Store(Snapshot{Generation: 2})
+	second <- ChangeMessage{Payload: `{"generation":2,"resource":"model"}`}
+	waitUntil(t, time.Second, func() bool { return manager.Generation() == 2 })
+
+	cancel()
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("manager did not stop")
+	}
+}
+
+func TestRedisCanBeDisabled(t *testing.T) {
+	store := newFakeManagerStore(3)
+	manager := newTestManager(store, slog.New(slog.NewTextHandler(&bytes.Buffer{}, nil)), 10*time.Millisecond)
+	if err := manager.AfterMutation(context.Background(), 3, "tenant", "tenant-1"); err != nil {
+		t.Fatal(err)
+	}
+	if manager.RedisConfigured() || manager.RedisConnected() {
+		t.Fatal("Redis unexpectedly reported as available")
+	}
+	if store.publishCalls.Load() != 0 || store.subscribeCalls.Load() != 0 {
+		t.Fatal("Redis operations were attempted while disabled")
+	}
+}
+
+func waitUntil(t *testing.T, timeout time.Duration, condition func() bool) {
+	t.Helper()
+	deadline := time.Now().Add(timeout)
+	for time.Now().Before(deadline) {
+		if condition() {
+			return
+		}
+		time.Sleep(5 * time.Millisecond)
+	}
+	t.Fatal("condition was not met before timeout")
+}