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