1package provider
2
3import (
4 "context"
5 "io"
6 "net/http"
7 "net/http/httptest"
8 "strings"
9 "testing"
10)
11
12type bedrockDiscoveryTestTransport func(*http.Request) (*http.Response, error)
13
14func (f bedrockDiscoveryTestTransport) RoundTrip(r *http.Request) (*http.Response, error) {
15 return f(r)
16}
17
18func TestDiscoverBedrockUsesSigV4WithBearer(t *testing.T) {
19 t.Setenv("AWS_BEARER_TOKEN_BEDROCK", "bearer")
20 t.Setenv("AWS_ACCESS_KEY_ID", "AKID")
21 t.Setenv("AWS_SECRET_ACCESS_KEY", "SECRET")
22 original := http.DefaultTransport
23 t.Cleanup(func() { http.DefaultTransport = original })
24 http.DefaultTransport = bedrockDiscoveryTestTransport(func(r *http.Request) (*http.Response, error) {
25 if !strings.Contains(r.Header.Get("Authorization"), "Credential=AKID/") {
26 t.Errorf("discovery request was not signed with SigV4")
27 }
28 body := `{"modelSummaries":[{"modelId":"test.new-model","outputModalities":["TEXT"]}]}`
29 if r.URL.Path == "/inference-profiles" {
30 body = `{"inferenceProfileSummaries":[]}`
31 }
32 return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body)), Header: make(http.Header)}, nil
33 })
34 models, err := DiscoverBedrock(context.Background(), "us-east-1")
35 if err != nil || len(models) != 1 || models[0].ID != "test.new-model" {
36 t.Fatalf("discovered models = %v, err = %v", models, err)
37 }
38 if bearer, sigv4 := resolveBedrockAuth("bearer"); bearer != "bearer" || sigv4 != nil {
39 t.Fatal("inference must continue to prefer the bearer token")
40 }
41}
42
43func TestBedrockEmitsText(t *testing.T) {
44 if !bedrockEmitsText(nil) {
45 t.Error("nil modalities should default to text")
46 }
47 if !bedrockEmitsText([]string{"TEXT"}) {
48 t.Error("TEXT should emit text")
49 }
50 if bedrockEmitsText([]string{"IMAGE"}) {
51 t.Error("IMAGE-only should not emit text")
52 }
53 if !bedrockEmitsText([]string{"IMAGE", "TEXT"}) {
54 t.Error("mixed with TEXT should emit text")
55 }
56}
57
58func TestResolveBedrockAuthBearerVsSigV4(t *testing.T) {
59 t.Setenv("AWS_BEARER_TOKEN_BEDROCK", "")
60 t.Setenv("AWS_ACCESS_KEY_ID", "")
61 t.Setenv("AWS_SECRET_ACCESS_KEY", "")
62 t.Setenv("AWS_SESSION_TOKEN", "")
63 t.Setenv("AWS_PROFILE", "")
64
65 // Explicit bearer via apiKey.
66 if b, s := resolveBedrockAuth("tok-123"); b != "tok-123" || s != nil {
67 t.Errorf("explicit bearer: got (%q,%v)", b, s)
68 }
69
70 // Sentinel <aws> is not a bearer; falls through to SigV4 env.
71 t.Setenv("AWS_ACCESS_KEY_ID", "AKID")
72 t.Setenv("AWS_SECRET_ACCESS_KEY", "SECRET")
73 b, s := resolveBedrockAuth("<aws>")
74 if b != "" || s == nil {
75 t.Fatalf("sentinel should resolve SigV4, got (%q,%v)", b, s)
76 }
77 if s.accessKeyID != "AKID" || s.secretAccessKey != "SECRET" {
78 t.Errorf("sigv4 creds = %+v", s)
79 }
80}
81
82// TestBedrockListModelIDs parses control-plane responses shaped exactly
83// like the real ones captured from a live account (see the doc note in
84// discover_bedrock.go): ListFoundationModels returns modelSummaries with
85// modelId + outputModalities; ListInferenceProfiles returns
86// inferenceProfileSummaries with inferenceProfileId and paginates via
87// nextToken. Embedding/image-only models are filtered out; profile IDs
88// are included; duplicates are de-duplicated.
89func TestBedrockListModelIDs(t *testing.T) {
90 page1 := `{"inferenceProfileSummaries":[{"inferenceProfileId":"us.anthropic.claude-opus-5-5"}],"nextToken":"N+E/T= &"}`
91 page2 := `{"inferenceProfileSummaries":[{"inferenceProfileId":"eu.anthropic.claude-sonnet-5"}]}`
92 foundation := `{"modelSummaries":[
93 {"modelId":"anthropic.claude-opus-5-5","outputModalities":["TEXT"]},
94 {"modelId":"amazon.titan-embed-text-v2","outputModalities":["EMBEDDING"]},
95 {"modelId":"stability.sd3-large","outputModalities":["IMAGE"]},
96 {"modelId":"anthropic.claude-opus-5-5","outputModalities":["TEXT"]}
97 ]}`
98
99 var ipHits int
100 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
101 w.Header().Set("content-type", "application/json")
102 switch {
103 case strings.Contains(r.URL.Path, "/foundation-models"):
104 w.Write([]byte(foundation))
105 case strings.Contains(r.URL.Path, "/inference-profiles"):
106 ipHits++
107 switch token := r.URL.Query().Get("nextToken"); token {
108 case "":
109 w.Write([]byte(page1))
110 case "N+E/T= &":
111 w.Write([]byte(page2))
112 default:
113 t.Errorf("unexpected pagination token %q", token)
114 w.WriteHeader(http.StatusBadRequest)
115 }
116 default:
117 w.WriteHeader(http.StatusNotFound)
118 }
119 }))
120 defer srv.Close()
121
122 creds := &bedrockSigV4Creds{accessKeyID: "AKID", secretAccessKey: "SECRET"}
123 ids, err := bedrockListModelIDs(context.Background(), srv.Client(), creds, "us-east-1", srv.URL)
124 if err != nil {
125 t.Fatalf("bedrockListModelIDs: %v", err)
126 }
127
128 want := []string{
129 "anthropic.claude-opus-5-5", // text foundation model, de-duplicated
130 "us.anthropic.claude-opus-5-5", // profile page 1
131 "eu.anthropic.claude-sonnet-5", // profile page 2 (pagination followed)
132 }
133 got := map[string]int{}
134 for _, id := range ids {
135 got[id]++
136 }
137 for _, w := range want {
138 if got[w] != 1 {
139 t.Errorf("expected exactly one %q, got %d (all=%v)", w, got[w], ids)
140 }
141 }
142 // Embedding and image-only models must be filtered out.
143 for _, bad := range []string{"amazon.titan-embed-text-v2", "stability.sd3-large"} {
144 if got[bad] != 0 {
145 t.Errorf("non-text model %q should be filtered out", bad)
146 }
147 }
148 if ipHits != 2 {
149 t.Errorf("expected 2 inference-profile page requests (pagination), got %d", ipHits)
150 }
151}
152
153// TestBedrockControlPlaneHTTPError confirms non-2xx responses surface as
154// errors with the status and body, so discovery fails loudly (and the
155// caller then falls back to the catalog) rather than silently parsing junk.
156func TestBedrockControlPlaneHTTPError(t *testing.T) {
157 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
158 w.WriteHeader(http.StatusForbidden)
159 w.Write([]byte(`{"message":"not authorized"}`))
160 }))
161 defer srv.Close()
162
163 creds := &bedrockSigV4Creds{accessKeyID: "AKID", secretAccessKey: "SECRET"}
164 _, err := bedrockSignedControlPlane(context.Background(), srv.Client(), creds, "us-east-1", http.MethodGet, srv.URL+"/foundation-models", nil)
165 if err == nil {
166 t.Fatal("expected error on 403")
167 }
168 if !strings.Contains(err.Error(), "403") || !strings.Contains(err.Error(), "not authorized") {
169 t.Errorf("error should carry status and body: %v", err)
170 }
171}