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`) }