xianfeng wang
bbca507792
* discovery * discovery * wk discovery * wk discovery * wk discovery * wk discovery * wk discovery * wk discovery * discovery example and default * position param * Fix test * dev auto login * fix test * fix conflict
255 lines
7.9 KiB
Go
255 lines
7.9 KiB
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":
|
|
_ = 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":
|
|
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)
|
|
}
|
|
}
|