| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380 |
- //go:build unit
- package service
- import (
- "context"
- "encoding/json"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/Wei-Shaw/sub2api/internal/pkg/websearch"
- "github.com/stretchr/testify/require"
- )
- // --- isOnlyWebSearchToolInBody ---
- func TestIsOnlyWebSearchToolInBody_WebSearchType(t *testing.T) {
- require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"type":"web_search"}]}`)))
- }
- func TestIsOnlyWebSearchToolInBody_WebSearch2025Type(t *testing.T) {
- require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"type":"web_search_20250305"}]}`)))
- }
- func TestIsOnlyWebSearchToolInBody_GoogleSearchType(t *testing.T) {
- require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"type":"google_search"}]}`)))
- }
- func TestIsOnlyWebSearchToolInBody_NameWebSearch(t *testing.T) {
- require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"name":"web_search"}]}`)))
- }
- func TestIsOnlyWebSearchToolInBody_NameWebSearch2025(t *testing.T) {
- require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"name":"web_search_20250305"}]}`)))
- }
- func TestIsOnlyWebSearchToolInBody_NameGoogleSearch(t *testing.T) {
- require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"name":"google_search"}]}`)))
- }
- func TestIsOnlyWebSearchToolInBody_MultipleTools(t *testing.T) {
- require.False(t, isOnlyWebSearchToolInBody(
- []byte(`{"tools":[{"type":"web_search"},{"type":"text_editor"}]}`)))
- }
- func TestIsOnlyWebSearchToolInBody_NoTools(t *testing.T) {
- require.False(t, isOnlyWebSearchToolInBody([]byte(`{"model":"claude-3"}`)))
- }
- func TestIsOnlyWebSearchToolInBody_EmptyToolsArray(t *testing.T) {
- require.False(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[]}`)))
- }
- func TestIsOnlyWebSearchToolInBody_NonWebSearchTool(t *testing.T) {
- require.False(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"type":"text_editor"}]}`)))
- }
- func TestIsOnlyWebSearchToolInBody_ToolsNotArray(t *testing.T) {
- require.False(t, isOnlyWebSearchToolInBody([]byte(`{"tools":"web_search"}`)))
- }
- // --- extractSearchQueryFromBody ---
- func TestExtractSearchQueryFromBody_StringContent(t *testing.T) {
- body := `{"messages":[{"role":"user","content":"what is golang"}]}`
- require.Equal(t, "what is golang", extractSearchQueryFromBody([]byte(body)))
- }
- func TestExtractSearchQueryFromBody_ArrayContent(t *testing.T) {
- body := `{"messages":[{"role":"user","content":[{"type":"text","text":"search this"}]}]}`
- require.Equal(t, "search this", extractSearchQueryFromBody([]byte(body)))
- }
- func TestExtractSearchQueryFromBody_MultipleMessages(t *testing.T) {
- body := `{"messages":[{"role":"user","content":"first"},{"role":"assistant","content":"ok"},{"role":"user","content":"second"}]}`
- require.Equal(t, "second", extractSearchQueryFromBody([]byte(body)))
- }
- func TestExtractSearchQueryFromBody_LastMessageNotUser(t *testing.T) {
- body := `{"messages":[{"role":"user","content":"q"},{"role":"assistant","content":"a"}]}`
- require.Equal(t, "", extractSearchQueryFromBody([]byte(body)))
- }
- func TestExtractSearchQueryFromBody_EmptyMessages(t *testing.T) {
- require.Equal(t, "", extractSearchQueryFromBody([]byte(`{"messages":[]}`)))
- }
- func TestExtractSearchQueryFromBody_NoMessages(t *testing.T) {
- require.Equal(t, "", extractSearchQueryFromBody([]byte(`{"model":"claude-3"}`)))
- }
- func TestExtractSearchQueryFromBody_ArrayContentSkipsEmptyText(t *testing.T) {
- body := `{"messages":[{"role":"user","content":[{"type":"image"},{"type":"text","text":""},{"type":"text","text":"real query"}]}]}`
- require.Equal(t, "real query", extractSearchQueryFromBody([]byte(body)))
- }
- func TestExtractSearchQueryFromBody_ArrayContentNoTextBlock(t *testing.T) {
- body := `{"messages":[{"role":"user","content":[{"type":"image","source":{}}]}]}`
- require.Equal(t, "", extractSearchQueryFromBody([]byte(body)))
- }
- // --- buildSearchResultBlocks ---
- func TestBuildSearchResultBlocks_WithResults(t *testing.T) {
- results := []websearch.SearchResult{
- {URL: "https://a.com", Title: "A", Snippet: "snippet a", PageAge: "2 days"},
- {URL: "https://b.com", Title: "B", Snippet: "snippet b"},
- }
- blocks := buildSearchResultBlocks(results)
- require.Len(t, blocks, 2)
- require.Equal(t, "web_search_result", blocks[0]["type"])
- require.Equal(t, "https://a.com", blocks[0]["url"])
- require.Equal(t, "snippet a", blocks[0]["page_content"])
- require.Equal(t, "2 days", blocks[0]["page_age"])
- // Second result has no PageAge
- require.Equal(t, "https://b.com", blocks[1]["url"])
- _, hasPageAge := blocks[1]["page_age"]
- require.False(t, hasPageAge)
- }
- func TestBuildSearchResultBlocks_Empty(t *testing.T) {
- blocks := buildSearchResultBlocks(nil)
- require.Empty(t, blocks)
- }
- func TestBuildSearchResultBlocks_SnippetEmpty(t *testing.T) {
- blocks := buildSearchResultBlocks([]websearch.SearchResult{{URL: "https://x.com", Title: "X", Snippet: ""}})
- _, hasContent := blocks[0]["page_content"]
- require.False(t, hasContent)
- }
- // --- buildTextSummary ---
- func TestBuildTextSummary_WithResults(t *testing.T) {
- results := []websearch.SearchResult{
- {URL: "https://a.com", Title: "A", Snippet: "desc a"},
- }
- summary := buildTextSummary("test query", results)
- require.Contains(t, summary, "test query")
- require.Contains(t, summary, "1. **A**")
- require.Contains(t, summary, "https://a.com")
- }
- func TestBuildTextSummary_NoResults(t *testing.T) {
- summary := buildTextSummary("test", nil)
- require.Contains(t, summary, "No search results found for: test")
- }
- // --- shouldEmulateWebSearch ---
- // webSearchToolBody is a valid request body with exactly one web_search tool.
- var webSearchToolBody = []byte(`{"tools":[{"type":"web_search"}],"messages":[{"role":"user","content":"test"}]}`)
- // nonWebSearchToolBody is a request body without web_search tool.
- var nonWebSearchToolBody = []byte(`{"tools":[{"type":"text_editor"}],"messages":[{"role":"user","content":"test"}]}`)
- // newAnthropicAPIKeyAccount creates a test Account with the given web search emulation mode.
- func newAnthropicAPIKeyAccount(mode string) *Account {
- return &Account{
- ID: 1,
- Platform: PlatformAnthropic,
- Type: AccountTypeAPIKey,
- Extra: map[string]any{featureKeyWebSearchEmulation: mode},
- }
- }
- // setGlobalWebSearchConfig stores a config in the global cache used by SettingService.IsWebSearchEmulationEnabled.
- func setGlobalWebSearchConfig(cfg *WebSearchEmulationConfig) {
- webSearchEmulationCache.Store(&cachedWebSearchEmulationConfig{
- config: cfg,
- expiresAt: time.Now().Add(10 * time.Minute).UnixNano(),
- })
- }
- // clearGlobalWebSearchConfig resets the global cache to force re-read.
- func clearGlobalWebSearchConfig() {
- webSearchEmulationCache.Store((*cachedWebSearchEmulationConfig)(nil))
- }
- // newSettingServiceForWebSearchTest creates a SettingService with a mock repo pre-loaded with config.
- func newSettingServiceForWebSearchTest(enabled bool) *SettingService {
- repo := newMockSettingRepo()
- cfg := &WebSearchEmulationConfig{
- Enabled: enabled,
- Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "sk-test"}},
- }
- data, _ := json.Marshal(cfg)
- repo.data[SettingKeyWebSearchEmulationConfig] = string(data)
- return NewSettingService(repo, &config.Config{})
- }
- // newChannelServiceWithCache creates a ChannelService with a pre-built cache containing the channel.
- func newChannelServiceWithCache(groupID int64, ch *Channel) *ChannelService {
- svc := &ChannelService{}
- cache := &channelCache{
- channelByGroupID: map[int64]*Channel{groupID: ch},
- byID: map[int64]*Channel{ch.ID: ch},
- groupPlatform: map[int64]string{},
- loadedAt: time.Now(),
- }
- svc.cache.Store(cache)
- return svc
- }
- func TestShouldEmulateWebSearch_NilManager(t *testing.T) {
- SetWebSearchManager(nil)
- defer SetWebSearchManager(nil)
- settingSvc := newSettingServiceForWebSearchTest(true)
- setGlobalWebSearchConfig(&WebSearchEmulationConfig{
- Enabled: true,
- Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
- })
- defer clearGlobalWebSearchConfig()
- svc := &GatewayService{settingService: settingSvc}
- account := newAnthropicAPIKeyAccount(WebSearchModeEnabled)
- require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, nil, webSearchToolBody))
- }
- func TestShouldEmulateWebSearch_NotOnlyWebSearchTool(t *testing.T) {
- mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
- SetWebSearchManager(mgr)
- defer SetWebSearchManager(nil)
- settingSvc := newSettingServiceForWebSearchTest(true)
- setGlobalWebSearchConfig(&WebSearchEmulationConfig{
- Enabled: true,
- Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
- })
- defer clearGlobalWebSearchConfig()
- svc := &GatewayService{settingService: settingSvc}
- account := newAnthropicAPIKeyAccount(WebSearchModeEnabled)
- require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, nil, nonWebSearchToolBody))
- }
- func TestShouldEmulateWebSearch_GlobalDisabled(t *testing.T) {
- mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
- SetWebSearchManager(mgr)
- defer SetWebSearchManager(nil)
- // Global config disabled
- setGlobalWebSearchConfig(&WebSearchEmulationConfig{
- Enabled: false,
- Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
- })
- defer clearGlobalWebSearchConfig()
- settingSvc := newSettingServiceForWebSearchTest(false)
- svc := &GatewayService{settingService: settingSvc}
- account := newAnthropicAPIKeyAccount(WebSearchModeEnabled)
- require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, nil, webSearchToolBody))
- }
- func TestShouldEmulateWebSearch_AccountDisabled(t *testing.T) {
- mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
- SetWebSearchManager(mgr)
- defer SetWebSearchManager(nil)
- setGlobalWebSearchConfig(&WebSearchEmulationConfig{
- Enabled: true,
- Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
- })
- defer clearGlobalWebSearchConfig()
- settingSvc := newSettingServiceForWebSearchTest(true)
- svc := &GatewayService{settingService: settingSvc}
- account := newAnthropicAPIKeyAccount(WebSearchModeDisabled)
- require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, nil, webSearchToolBody))
- }
- func TestShouldEmulateWebSearch_AccountEnabled(t *testing.T) {
- mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
- SetWebSearchManager(mgr)
- defer SetWebSearchManager(nil)
- setGlobalWebSearchConfig(&WebSearchEmulationConfig{
- Enabled: true,
- Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
- })
- defer clearGlobalWebSearchConfig()
- settingSvc := newSettingServiceForWebSearchTest(true)
- svc := &GatewayService{settingService: settingSvc}
- account := newAnthropicAPIKeyAccount(WebSearchModeEnabled)
- require.True(t, svc.shouldEmulateWebSearch(context.Background(), account, nil, webSearchToolBody))
- }
- func TestShouldEmulateWebSearch_DefaultMode_ChannelEnabled(t *testing.T) {
- mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
- SetWebSearchManager(mgr)
- defer SetWebSearchManager(nil)
- setGlobalWebSearchConfig(&WebSearchEmulationConfig{
- Enabled: true,
- Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
- })
- defer clearGlobalWebSearchConfig()
- settingSvc := newSettingServiceForWebSearchTest(true)
- ch := &Channel{
- ID: 10,
- Status: StatusActive,
- FeaturesConfig: map[string]any{
- featureKeyWebSearchEmulation: map[string]any{PlatformAnthropic: true},
- },
- }
- channelSvc := newChannelServiceWithCache(42, ch)
- svc := &GatewayService{settingService: settingSvc, channelService: channelSvc}
- account := newAnthropicAPIKeyAccount(WebSearchModeDefault)
- groupID := int64(42)
- require.True(t, svc.shouldEmulateWebSearch(context.Background(), account, &groupID, webSearchToolBody))
- }
- func TestShouldEmulateWebSearch_DefaultMode_ChannelDisabled(t *testing.T) {
- mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
- SetWebSearchManager(mgr)
- defer SetWebSearchManager(nil)
- setGlobalWebSearchConfig(&WebSearchEmulationConfig{
- Enabled: true,
- Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
- })
- defer clearGlobalWebSearchConfig()
- settingSvc := newSettingServiceForWebSearchTest(true)
- ch := &Channel{
- ID: 10,
- Status: StatusActive,
- FeaturesConfig: map[string]any{
- featureKeyWebSearchEmulation: map[string]any{PlatformAnthropic: false},
- },
- }
- channelSvc := newChannelServiceWithCache(42, ch)
- svc := &GatewayService{settingService: settingSvc, channelService: channelSvc}
- account := newAnthropicAPIKeyAccount(WebSearchModeDefault)
- groupID := int64(42)
- require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, &groupID, webSearchToolBody))
- }
- func TestShouldEmulateWebSearch_DefaultMode_NilGroupID(t *testing.T) {
- mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
- SetWebSearchManager(mgr)
- defer SetWebSearchManager(nil)
- setGlobalWebSearchConfig(&WebSearchEmulationConfig{
- Enabled: true,
- Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
- })
- defer clearGlobalWebSearchConfig()
- settingSvc := newSettingServiceForWebSearchTest(true)
- svc := &GatewayService{settingService: settingSvc}
- account := newAnthropicAPIKeyAccount(WebSearchModeDefault)
- // nil groupID + default mode → falls through to channel check → returns false
- require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, nil, webSearchToolBody))
- }
- func TestShouldEmulateWebSearch_DefaultMode_NilChannelService(t *testing.T) {
- mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
- SetWebSearchManager(mgr)
- defer SetWebSearchManager(nil)
- setGlobalWebSearchConfig(&WebSearchEmulationConfig{
- Enabled: true,
- Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
- })
- defer clearGlobalWebSearchConfig()
- settingSvc := newSettingServiceForWebSearchTest(true)
- svc := &GatewayService{settingService: settingSvc, channelService: nil}
- account := newAnthropicAPIKeyAccount(WebSearchModeDefault)
- groupID := int64(42)
- // nil channelService + default mode → returns false
- require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, &groupID, webSearchToolBody))
- }
|