1package bedrock
2
3import (
4 "encoding/json"
5 "net/http"
6 "net/http/httptest"
7 "testing"
8
9 "charm.land/fantasy"
10 "github.com/stretchr/testify/assert"
11 "github.com/stretchr/testify/require"
12)
13
14func TestUserAgent(t *testing.T) {
15 t.Parallel()
16
17 newUAServer := func() (*httptest.Server, *[]map[string]string) {
18 var captured []map[string]string
19 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
20 h := make(map[string]string)
21 for k, v := range r.Header {
22 if len(v) > 0 {
23 h[k] = v[0]
24 }
25 }
26 captured = append(captured, h)
27
28 w.Header().Set("Content-Type", "application/json")
29 _ = json.NewEncoder(w).Encode(mockAnthropicResponse())
30 }))
31 return server, &captured
32 }
33
34 prompt := fantasy.Prompt{
35 {
36 Role: fantasy.MessageRoleUser,
37 Content: []fantasy.MessagePart{fantasy.TextPart{Text: "Hi"}},
38 },
39 }
40
41 t.Run("default UA applied", func(t *testing.T) {
42 t.Parallel()
43 server, captured := newUAServer()
44 defer server.Close()
45
46 p, err := New(
47 WithAPIKey("k"),
48 WithSkipAuth(true),
49 WithHTTPClient(&http.Client{Transport: redirectTransport(server.URL)}),
50 )
51 require.NoError(t, err)
52 model, _ := p.LanguageModel(t.Context(), "us.anthropic.claude-sonnet-4-20250514-v1:0")
53 _, _ = model.Generate(t.Context(), fantasy.Call{Prompt: prompt})
54
55 require.Len(t, *captured, 1)
56 assert.Equal(t, "Charm Fantasy/"+fantasy.Version, (*captured)[0]["User-Agent"])
57 })
58
59 t.Run("WithUserAgent wins over default", func(t *testing.T) {
60 t.Parallel()
61 server, captured := newUAServer()
62 defer server.Close()
63
64 p, err := New(
65 WithAPIKey("k"),
66 WithSkipAuth(true),
67 WithHTTPClient(&http.Client{Transport: redirectTransport(server.URL)}),
68 WithUserAgent("explicit-ua"),
69 )
70 require.NoError(t, err)
71 model, _ := p.LanguageModel(t.Context(), "us.anthropic.claude-sonnet-4-20250514-v1:0")
72 _, _ = model.Generate(t.Context(), fantasy.Call{Prompt: prompt})
73
74 require.Len(t, *captured, 1)
75 assert.Equal(t, "explicit-ua", (*captured)[0]["User-Agent"])
76 })
77
78 t.Run("WithUserAgent wins over WithHeaders", func(t *testing.T) {
79 t.Parallel()
80 server, captured := newUAServer()
81 defer server.Close()
82
83 p, err := New(
84 WithAPIKey("k"),
85 WithSkipAuth(true),
86 WithHTTPClient(&http.Client{Transport: redirectTransport(server.URL)}),
87 WithHeaders(map[string]string{"User-Agent": "from-headers"}),
88 WithUserAgent("explicit-ua"),
89 )
90 require.NoError(t, err)
91 model, _ := p.LanguageModel(t.Context(), "us.anthropic.claude-sonnet-4-20250514-v1:0")
92 _, _ = model.Generate(t.Context(), fantasy.Call{Prompt: prompt})
93
94 require.Len(t, *captured, 1)
95 assert.Equal(t, "explicit-ua", (*captured)[0]["User-Agent"])
96 })
97}
98
99type redirectRoundTripper struct {
100 target string
101}
102
103func redirectTransport(target string) *redirectRoundTripper {
104 return &redirectRoundTripper{target: target}
105}
106
107func (rt *redirectRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
108 req = req.Clone(req.Context())
109 req.URL.Scheme = "http"
110 req.URL.Host = (&redirectRoundTripper{target: rt.target}).host()
111 return http.DefaultTransport.RoundTrip(req)
112}
113
114func (rt *redirectRoundTripper) host() string {
115 u := rt.target
116 if len(u) > 7 && u[:7] == "http://" {
117 return u[7:]
118 }
119 if len(u) > 8 && u[:8] == "https://" {
120 return u[8:]
121 }
122 return u
123}
124
125func mockAnthropicResponse() map[string]any {
126 return map[string]any{
127 "id": "msg_01Test",
128 "type": "message",
129 "role": "assistant",
130 "model": "claude-sonnet-4-20250514",
131 "content": []any{
132 map[string]any{
133 "type": "text",
134 "text": "Hi there",
135 },
136 },
137 "stop_reason": "end_turn",
138 "stop_sequence": "",
139 "usage": map[string]any{
140 "cache_creation": map[string]any{
141 "ephemeral_1h_input_tokens": 0,
142 "ephemeral_5m_input_tokens": 0,
143 },
144 "cache_creation_input_tokens": 0,
145 "cache_read_input_tokens": 0,
146 "input_tokens": 5,
147 "output_tokens": 2,
148 "server_tool_use": map[string]any{
149 "web_search_requests": 0,
150 },
151 "service_tier": "standard",
152 },
153 }
154}