1package openrouter
2
3import (
4 "encoding/json"
5 "net/http"
6 "net/http/httptest"
7 "testing"
8
9 "charm.land/fantasy"
10 "charm.land/fantasy/providers/openai"
11 "github.com/stretchr/testify/assert"
12 "github.com/stretchr/testify/require"
13)
14
15func TestUserAgent(t *testing.T) {
16 t.Parallel()
17
18 newUAServer := func() (*httptest.Server, *[]map[string]string) {
19 var captured []map[string]string
20 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
21 h := make(map[string]string)
22 for k, v := range r.Header {
23 if len(v) > 0 {
24 h[k] = v[0]
25 }
26 }
27 captured = append(captured, h)
28
29 w.Header().Set("Content-Type", "application/json")
30 _ = json.NewEncoder(w).Encode(mockOpenAIResponse())
31 }))
32 return server, &captured
33 }
34
35 withBaseURL := func(url string) Option {
36 return func(o *options) {
37 o.openaiOptions = append(o.openaiOptions, openai.WithBaseURL(url))
38 }
39 }
40
41 prompt := fantasy.Prompt{
42 {
43 Role: fantasy.MessageRoleUser,
44 Content: []fantasy.MessagePart{fantasy.TextPart{Text: "Hi"}},
45 },
46 }
47
48 t.Run("default UA applied", func(t *testing.T) {
49 t.Parallel()
50 server, captured := newUAServer()
51 defer server.Close()
52
53 p, err := New(WithAPIKey("k"), withBaseURL(server.URL))
54 require.NoError(t, err)
55 model, _ := p.LanguageModel(t.Context(), "openai/gpt-4")
56 _, _ = model.Generate(t.Context(), fantasy.Call{Prompt: prompt})
57
58 require.Len(t, *captured, 1)
59 assert.Equal(t, "Charm Fantasy/"+fantasy.Version, (*captured)[0]["User-Agent"])
60 })
61
62 t.Run("WithUserAgent wins over default", func(t *testing.T) {
63 t.Parallel()
64 server, captured := newUAServer()
65 defer server.Close()
66
67 p, err := New(WithAPIKey("k"), withBaseURL(server.URL), WithUserAgent("explicit-ua"))
68 require.NoError(t, err)
69 model, _ := p.LanguageModel(t.Context(), "openai/gpt-4")
70 _, _ = model.Generate(t.Context(), fantasy.Call{Prompt: prompt})
71
72 require.Len(t, *captured, 1)
73 assert.Equal(t, "explicit-ua", (*captured)[0]["User-Agent"])
74 })
75
76 t.Run("WithUserAgent wins over WithHeaders", func(t *testing.T) {
77 t.Parallel()
78 server, captured := newUAServer()
79 defer server.Close()
80
81 p, err := New(
82 WithAPIKey("k"),
83 withBaseURL(server.URL),
84 WithHeaders(map[string]string{"User-Agent": "from-headers"}),
85 WithUserAgent("explicit-ua"),
86 )
87 require.NoError(t, err)
88 model, _ := p.LanguageModel(t.Context(), "openai/gpt-4")
89 _, _ = model.Generate(t.Context(), fantasy.Call{Prompt: prompt})
90
91 require.Len(t, *captured, 1)
92 assert.Equal(t, "explicit-ua", (*captured)[0]["User-Agent"])
93 })
94}
95
96func mockOpenAIResponse() map[string]any {
97 return map[string]any{
98 "id": "chatcmpl-test",
99 "object": "chat.completion",
100 "created": 1711115037,
101 "model": "openai/gpt-4",
102 "choices": []map[string]any{
103 {
104 "index": 0,
105 "message": map[string]any{
106 "role": "assistant",
107 "content": "Hi there",
108 },
109 "finish_reason": "stop",
110 },
111 },
112 "usage": map[string]any{
113 "prompt_tokens": 4,
114 "total_tokens": 6,
115 "completion_tokens": 2,
116 },
117 }
118}