dingtalk-workspace-cli/internal/app/root_cache_refresh_test.go

package app

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"

	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/cache"
	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/cli"
	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/market"
	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/internal/transport"
	"github.com/DingTalk-Real-AI/dingtalk-workspace-cli/pkg/edition"
)

func TestCacheRefreshClearsExistingCachesAndSkipsCLISkippedServers(t *testing.T) {
	cacheDir := t.TempDir()
	t.Setenv(cli.CacheDirEnv, cacheDir)

	var skippedRuntimeCalls atomic.Int32

	var srv *httptest.Server
	srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/cli/discovery/apis/cedar":
			_ = json.NewEncoder(w).Encode(market.ListResponse{
				Metadata: market.ListMetadata{Count: 2},
				Servers: []market.ServerEnvelope{
					{
						Server: market.RegistryServer{
							Name: "Active Service",
							Remotes: []market.RegistryRemote{
								{Type: "streamable-http", URL: srv.URL + "/mcp/active"},
							},
						},
						Meta: market.EnvelopeMeta{
							Registry: market.RegistryMetadata{Status: "active"},
							CLI:      market.CLIOverlay{ID: "active", Command: "active"},
						},
					},
					{
						Server: market.RegistryServer{
							Name: "Skipped Service",
							Remotes: []market.RegistryRemote{
								{Type: "streamable-http", URL: srv.URL + "/mcp/skipped"},
							},
						},
						Meta: market.EnvelopeMeta{
							Registry: market.RegistryMetadata{Status: "active"},
							CLI:      market.CLIOverlay{ID: "legacy", Command: "legacy", Skip: true},
						},
					},
				},
			})
		case "/mcp/active":
			http.Error(w, "active runtime unavailable", http.StatusInternalServerError)
		case "/mcp/skipped":
			skippedRuntimeCalls.Add(1)
			http.Error(w, "skipped runtime should not be called", http.StatusInternalServerError)
		default:
			http.NotFound(w, r)
		}
	}))
	defer srv.Close()

	store := cache.NewStore(cacheDir)
	const partition = "default/default"
	activeKey := market.ServerKey(srv.URL + "/mcp/active")
	skippedKey := market.ServerKey(srv.URL + "/mcp/skipped")

	saveCachedRuntimeAndDetail(t, store, partition, activeKey)
	saveCachedRuntimeAndDetail(t, store, partition, skippedKey)
	saveCLIIDDetail(t, store, partition, "active")

	SetDiscoveryBaseURL(srv.URL)
	t.Cleanup(func() { SetDiscoveryBaseURL("") })

	cmd := newCacheCommand()
	var out bytes.Buffer
	cmd.SetOut(&out)
	cmd.SetErr(&out)
	cmd.SetArgs([]string{"refresh"})

	if err := cmd.Execute(); err != nil {
		t.Fatalf("Execute() error = %v", err)
	}

	if _, _, err := store.LoadTools(partition, activeKey); err == nil {
		t.Fatal("LoadTools(active) error = nil, want cache cleared before failed refresh")
	}
	if _, _, err := store.LoadDetail(partition, activeKey); err == nil {
		t.Fatal("LoadDetail(active) error = nil, want detail cache cleared before failed refresh")
	}
	if _, _, err := store.LoadDetail(partition, "active"); err != nil {
		t.Fatalf("LoadDetail(active CLI.ID) error = %v, want CLI metadata preserved on failed refresh", err)
	}
	if _, _, err := store.LoadTools(partition, skippedKey); err == nil {
		t.Fatal("LoadTools(skipped) error = nil, want skipped service cache removed")
	}
	if _, _, err := store.LoadDetail(partition, skippedKey); err == nil {
		t.Fatal("LoadDetail(skipped) error = nil, want skipped service detail cache removed")
	}
	if got := skippedRuntimeCalls.Load(); got != 0 {
		t.Fatalf("skipped runtime calls = %d, want 0", got)
	}
}

// TestCacheRefreshHonorsEditionDiscoveryURL asserts the `dws cache refresh`
// command routes its server-list fetch through edition.Hooks.DiscoveryURL /
// DiscoveryHeaders when they are set, instead of the default Market endpoint.
// Kept deliberately generic (no edition-specific strings) — concrete values
// belong to the overlay repo that installs the hooks, not to this open core.
func TestCacheRefreshHonorsEditionDiscoveryURL(t *testing.T) {
	cacheDir := t.TempDir()
	t.Setenv(cli.CacheDirEnv, cacheDir)

	var (
		editionHits atomic.Int32
		marketHits  atomic.Int32
		gotHeaders  atomic.Value // map[string]string
	)

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/cli/edition/apis":
			editionHits.Add(1)
			snapshot := map[string]string{
				"x-test-edition": r.Header.Get("x-test-edition"),
				"x-test-client":  r.Header.Get("x-test-client"),
			}
			gotHeaders.Store(snapshot)
			_ = json.NewEncoder(w).Encode(market.ListResponse{
				Metadata: market.ListMetadata{Count: 1},
				Servers: []market.ServerEnvelope{
					{
						Server: market.RegistryServer{
							Name:    "Edition Service",
							Remotes: []market.RegistryRemote{{Type: "streamable-http", URL: "https://example.invalid/mcp"}},
						},
						Meta: market.EnvelopeMeta{
							Registry: market.RegistryMetadata{Status: "active"},
							CLI:      market.CLIOverlay{ID: "edition-service", Command: "edition-service"},
						},
					},
				},
			})
		case "/cli/discovery/apis/cedar":
			marketHits.Add(1)
			http.Error(w, "market endpoint must not be called when edition DiscoveryURL is set", http.StatusNotFound)
		default:
			http.NotFound(w, r)
		}
	}))
	defer srv.Close()

	edition.Override(&edition.Hooks{
		Name:         "testing",
		DiscoveryURL: srv.URL + "/cli/edition/apis",
		DiscoveryHeaders: func() map[string]string {
			return map[string]string{
				"x-test-edition": "custom",
				"x-test-client":  "cli-refresh",
			}
		},
	})
	t.Cleanup(func() { edition.Override(&edition.Hooks{}) })

	SetDiscoveryBaseURL(srv.URL)
	t.Cleanup(func() { SetDiscoveryBaseURL("") })

	cmd := newCacheCommand()
	var out bytes.Buffer
	cmd.SetOut(&out)
	cmd.SetErr(&out)
	cmd.SetArgs([]string{"refresh"})

	if err := cmd.Execute(); err != nil {
		t.Fatalf("Execute() error = %v", err)
	}

	if got := editionHits.Load(); got != 1 {
		t.Fatalf("edition DiscoveryURL hits = %d, want 1", got)
	}
	if got := marketHits.Load(); got != 0 {
		t.Fatalf("market endpoint hits = %d, want 0 (edition DiscoveryURL must take precedence)", got)
	}

	headers, _ := gotHeaders.Load().(map[string]string)
	if headers == nil {
		t.Fatal("captured request headers = nil, want edition DiscoveryHeaders to be applied")
	}
	if headers["x-test-edition"] != "custom" {
		t.Fatalf("x-test-edition header = %q, want %q", headers["x-test-edition"], "custom")
	}
	if headers["x-test-client"] != "cli-refresh" {
		t.Fatalf("x-test-client header = %q, want %q", headers["x-test-client"], "cli-refresh")
	}
}

func saveCLIIDDetail(t *testing.T, store *cache.Store, partition, cliID string) {
	t.Helper()

	payload, err := json.Marshal(market.DetailResponse{
		Success: true,
		Result: market.DetailResult{
			Tools: []market.DetailTool{
				{ToolName: "stale_tool", ToolTitle: "Stale Tool"},
			},
		},
	})
	if err != nil {
		t.Fatalf("json.Marshal(cli detail payload) error = %v", err)
	}
	if err := store.SaveDetail(partition, cliID, cache.DetailSnapshot{
		MCPID:   0,
		Payload: payload,
	}); err != nil {
		t.Fatalf("SaveDetail(%s) error = %v", cliID, err)
	}
}

func saveCachedRuntimeAndDetail(t *testing.T, store *cache.Store, partition, serverKey string) {
	t.Helper()

	if err := store.SaveTools(partition, serverKey, cache.ToolsSnapshot{
		ServerKey:       serverKey,
		ProtocolVersion: "2025-03-26",
		Tools: []transport.ToolDescriptor{
			{Name: "stale_tool", Title: "Stale Tool"},
		},
	}); err != nil {
		t.Fatalf("SaveTools(%s) error = %v", serverKey, err)
	}

	payload, err := json.Marshal(market.DetailResponse{
		Success: true,
		Result: market.DetailResult{
			Tools: []market.DetailTool{
				{ToolName: "stale_tool", ToolTitle: "Stale Tool"},
			},
		},
	})
	if err != nil {
		t.Fatalf("json.Marshal(detail payload) error = %v", err)
	}
	if err := store.SaveDetail(partition, serverKey, cache.DetailSnapshot{
		MCPID:   0,
		Payload: payload,
	}); err != nil {
		t.Fatalf("SaveDetail(%s) error = %v", serverKey, err)
	}
}