sub2api/backend/internal/pkg/lspool/worker_manager_test.go

package lspool

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/network"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/stretchr/testify/require"
)

type fakeDockerClient struct {
	mu sync.Mutex

	listResp       []container.Summary
	listCalls      int
	createCalls    int
	startCalls     int
	stopCalls      int
	removeCalls    int
	inspectCalls   int
	removedIDs     []string
	createdConfigs []*container.Config
	inspectResp    container.InspectResponse
}

func (f *fakeDockerClient) ContainerList(ctx context.Context, options container.ListOptions) ([]container.Summary, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.listCalls++
	return append([]container.Summary(nil), f.listResp...), nil
}

func (f *fakeDockerClient) ContainerCreate(ctx context.Context, cfg *container.Config, hostConfig *container.HostConfig, networkingConfig *network.NetworkingConfig, platform *ocispec.Platform, containerName string) (container.CreateResponse, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.createCalls++
	f.createdConfigs = append(f.createdConfigs, cfg)
	return container.CreateResponse{ID: "worker-created"}, nil
}

func (f *fakeDockerClient) ContainerStart(ctx context.Context, containerID string, options container.StartOptions) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.startCalls++
	return nil
}

func (f *fakeDockerClient) ContainerInspect(ctx context.Context, containerID string) (container.InspectResponse, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.inspectCalls++
	return f.inspectResp, nil
}

func (f *fakeDockerClient) ContainerStop(ctx context.Context, containerID string, options container.StopOptions) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.stopCalls++
	return nil
}

func (f *fakeDockerClient) ContainerRemove(ctx context.Context, containerID string, options container.RemoveOptions) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.removeCalls++
	f.removedIDs = append(f.removedIDs, containerID)
	return nil
}

func (f *fakeDockerClient) Close() error { return nil }

func TestResolveWorkerProxyRejectsHTTP(t *testing.T) {
	_, _, err := resolveWorkerProxy("http://127.0.0.1:7890")
	require.Error(t, err)
	require.Contains(t, err.Error(), "only supports socks5/socks5h")
}

func TestProxyHashUsesNormalizedProxy(t *testing.T) {
	normalized, _, err := resolveWorkerProxy("socks5://user:pass@127.0.0.1:1080")
	require.NoError(t, err)
	require.Equal(t, "socks5h://user:pass@127.0.0.1:1080", normalized)

	hash1 := proxyHash(normalized)
	hash2 := proxyHash("socks5h://user:pass@127.0.0.1:1080")
	require.Equal(t, hash1, hash2)
}

func TestWorkerManagerRequiresToken(t *testing.T) {
	fakeDocker := &fakeDockerClient{}
	manager, err := newWorkerManager(workerManagerConfig{
		Image:          "worker:latest",
		Network:        "sub2api-network",
		DockerSocket:   "unix:///var/run/docker.sock",
		IdleTTL:        time.Minute,
		MaxActive:      2,
		StartupTimeout: time.Second,
		RequestTimeout: time.Second,
	}, fakeDocker)
	require.NoError(t, err)
	defer manager.Close()

	_, err = manager.GetOrCreate("9", "rk-1", "socks5h://user:pass@127.0.0.1:1080")
	require.Error(t, err)
	require.Contains(t, err.Error(), "missing access token")
}

func TestWorkerManagerReusesExistingHandleAndDedupesStateSync(t *testing.T) {
	var mu sync.Mutex
	var healthCalls int
	var readyCalls int
	var stateCalls int
	var stateBodies [][]byte

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/healthz":
			mu.Lock()
			healthCalls++
			mu.Unlock()
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte("ok"))
		case "/readyz":
			mu.Lock()
			readyCalls++
			mu.Unlock()
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte("ready"))
		case "/account/state":
			body, _ := io.ReadAll(r.Body)
			mu.Lock()
			stateCalls++
			stateBodies = append(stateBodies, body)
			mu.Unlock()
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte("ok"))
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	fakeDocker := &fakeDockerClient{}
	manager, err := newWorkerManager(workerManagerConfig{
		Image:          "worker:latest",
		Network:        "sub2api-network",
		DockerSocket:   "unix:///var/run/docker.sock",
		IdleTTL:        time.Minute,
		MaxActive:      4,
		StartupTimeout: time.Second,
		RequestTimeout: time.Second,
	}, fakeDocker)
	require.NoError(t, err)
	defer manager.Close()

	accountID := "9"
	proxyURL := "socks5h://user:pass@127.0.0.1:1080"
	hash := proxyHash(proxyURL)
	key := buildWorkerKey(accountID, hash)

	manager.SetAccountToken(accountID, "ya29.test", "refresh", time.Now().Add(time.Hour))
	manager.mu.Lock()
	manager.workers[key] = &workerHandle{
		Key:         key,
		AccountID:   accountID,
		ProxyURL:    proxyURL,
		ProxyHash:   hash,
		ContainerID: "existing-worker",
		Container:   "sub2api-ls-9-test",
		Address:     strings.TrimPrefix(server.URL, "http://"),
		AuthToken:   "worker-token",
		LastUsed:    time.Now(),
	}
	manager.mu.Unlock()

	inst1, err := manager.GetOrCreate(accountID, "rk-1", proxyURL)
	require.NoError(t, err)
	require.True(t, inst1.remote)
	require.Equal(t, replicaSlotIndex("rk-1", parseLSReplicaCount()), inst1.Replica)

	inst2, err := manager.GetOrCreate(accountID, "rk-1", proxyURL)
	require.NoError(t, err)
	require.True(t, inst2.remote)

	mu.Lock()
	defer mu.Unlock()
	require.GreaterOrEqual(t, healthCalls, 2)
	require.GreaterOrEqual(t, readyCalls, 2)
	require.Equal(t, 1, stateCalls, "state sync should be skipped when the payload hash is unchanged")
	require.Len(t, stateBodies, 1)

	var synced workerAccountState
	require.NoError(t, json.Unmarshal(stateBodies[0], &synced))
	require.True(t, synced.HasToken)
	require.Equal(t, "ya29.test", synced.AccessToken)
}

func TestWorkerManagerMaxActiveStopsNewWorkerCreation(t *testing.T) {
	fakeDocker := &fakeDockerClient{}
	manager, err := newWorkerManager(workerManagerConfig{
		Image:          "worker:latest",
		Network:        "sub2api-network",
		DockerSocket:   "unix:///var/run/docker.sock",
		IdleTTL:        time.Minute,
		MaxActive:      1,
		StartupTimeout: time.Second,
		RequestTimeout: time.Second,
	}, fakeDocker)
	require.NoError(t, err)
	defer manager.Close()

	manager.SetAccountToken("9", "ya29.test", "refresh", time.Now().Add(time.Hour))
	manager.mu.Lock()
	manager.workers["existing"] = &workerHandle{ContainerID: "existing", Container: "existing", LastUsed: time.Now()}
	manager.mu.Unlock()

	_, err = manager.GetOrCreate("9", "rk-new", "socks5h://user:pass@127.0.0.1:1080")
	require.Error(t, err)
	require.Contains(t, err.Error(), "limit reached")
	require.Equal(t, 0, fakeDocker.createCalls)
}

func TestWorkerManagerReconcileRemovesManagedContainers(t *testing.T) {
	fakeDocker := &fakeDockerClient{
		listResp: []container.Summary{
			{
				ID:    "old-worker-1",
				Names: []string{"/sub2api-ls-9-deadbeef"},
			},
			{
				ID:    "old-worker-2",
				Names: []string{"/sub2api-ls-10-beadfeed"},
			},
		},
	}

	manager, err := newWorkerManager(workerManagerConfig{
		Image:          "worker:latest",
		Network:        "sub2api-network",
		DockerSocket:   "unix:///var/run/docker.sock",
		IdleTTL:        time.Minute,
		MaxActive:      4,
		StartupTimeout: time.Second,
		RequestTimeout: time.Second,
	}, fakeDocker)
	require.NoError(t, err)
	defer manager.Close()

	require.Equal(t, 1, fakeDocker.listCalls)
	require.ElementsMatch(t, []string{"old-worker-1", "old-worker-2"}, fakeDocker.removedIDs)
}

func TestFakeDockerClientImplementsFilterAwareList(t *testing.T) {
	fakeDocker := &fakeDockerClient{}
	_, err := fakeDocker.ContainerList(context.Background(), container.ListOptions{Filters: filters.NewArgs()})
	require.NoError(t, err)
}

func TestShouldWarnWorkerNotReadySuppressesModelMappingPending(t *testing.T) {
	require.False(t, shouldWarnWorkerNotReady(http.StatusServiceUnavailable, "worker model mapping not ready for replica 0"))
	require.True(t, shouldWarnWorkerNotReady(http.StatusServiceUnavailable, "worker access token not configured"))
	require.True(t, shouldWarnWorkerNotReady(http.StatusBadGateway, "upstream failed"))
}

func TestWorkerManagerWaitForWorkerReadyStopsOnModelMappingUnavailable(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, "/readyz", r.URL.Path)
		w.WriteHeader(http.StatusServiceUnavailable)
		_, _ = w.Write([]byte(`model mapping unavailable for replica 0: oauth2: "unauthorized_client" "Unauthorized"`))
	}))
	defer server.Close()

	manager, err := newWorkerManager(workerManagerConfig{
		Image:          "worker:latest",
		Network:        "sub2api-network",
		DockerSocket:   "unix:///var/run/docker.sock",
		IdleTTL:        time.Minute,
		MaxActive:      1,
		StartupTimeout: time.Second,
		RequestTimeout: time.Second,
	}, &fakeDockerClient{})
	require.NoError(t, err)
	defer manager.Close()

	handle := &workerHandle{
		Container: "sub2api-ls-test",
		Address:   strings.TrimPrefix(server.URL, "http://"),
		AuthToken: "worker-token",
	}

	err = manager.waitForWorkerReady(handle, "")
	require.Error(t, err)
	require.ErrorIs(t, err, errLSModelMapDenied)
}

func TestWorkerManagerWaitForWorkerReadyIncludesLastBodyOnTimeout(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, "/readyz", r.URL.Path)
		w.WriteHeader(http.StatusServiceUnavailable)
		_, _ = w.Write([]byte("worker model mapping not ready for replica 0\n"))
	}))
	defer server.Close()

	manager, err := newWorkerManager(workerManagerConfig{
		Image:          "worker:latest",
		Network:        "sub2api-network",
		DockerSocket:   "unix:///var/run/docker.sock",
		IdleTTL:        time.Minute,
		MaxActive:      1,
		StartupTimeout: 100 * time.Millisecond,
		RequestTimeout: time.Second,
	}, &fakeDockerClient{})
	require.NoError(t, err)
	defer manager.Close()

	handle := &workerHandle{
		Container: "sub2api-ls-test",
		Address:   strings.TrimPrefix(server.URL, "http://"),
		AuthToken: "worker-token",
	}

	err = manager.waitForWorkerReady(handle, "")
	require.Error(t, err)
	require.Contains(t, err.Error(), `last_status=503`)
	require.Contains(t, err.Error(), `last_body="worker model mapping not ready for replica 0`)
}