| 1 | package provider |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "reflect" |
| 7 | "testing" |
| 8 | ) |
| 9 | |
| 10 | type requestGateFixture struct { |
| 11 | blocked bool |
| 12 | ref string |
| 13 | err error |
| 14 | } |
| 15 | |
| 16 | func (g *requestGateFixture) BeforeModelRequest(ref string) error { |
| 17 | g.ref = ref |
| 18 | if g.blocked { |
| 19 | return errors.New("blocked") |
| 20 | } |
| 21 | return nil |
| 22 | } |
| 23 | func (g *requestGateFixture) ModelRequestFailed(ref string, err error) { |
| 24 | g.ref = ref |
| 25 | g.err = err |
| 26 | g.blocked = true |
| 27 | } |
| 28 | |
| 29 | type gatedProviderFixture struct { |
| 30 | calls int |
| 31 | request Request |
| 32 | streamError bool |
| 33 | } |
| 34 | |
| 35 | func (*gatedProviderFixture) Name() string { return "connection" } |
| 36 | func (*gatedProviderFixture) ModelInfo() ModelInfo { return ModelInfo{ID: "model"} } |
| 37 | func (p *gatedProviderFixture) Stream(_ context.Context, req Request) (<-chan Chunk, error) { |
| 38 | p.calls++ |
| 39 | p.request = req |
| 40 | err := &AuthError{Provider: "connection", Status: 403} |
| 41 | if !p.streamError { |
| 42 | return nil, err |
| 43 | } |
| 44 | ch := make(chan Chunk, 1) |
| 45 | ch <- Chunk{Type: ChunkError, Err: err} |
| 46 | close(ch) |
| 47 | return ch, nil |
| 48 | } |
| 49 | |
| 50 | func TestRequestGateObservesHTTPAndStreamRejectionsWithoutChangingRequest(t *testing.T) { |
| 51 | for _, streamError := range []bool{false, true} { |
| 52 | p := &gatedProviderFixture{streamError: streamError} |
| 53 | gate := &requestGateFixture{} |
| 54 | ctx, cancel := context.WithCancel(WithRequestGate(context.Background(), gate)) |
| 55 | req := Request{Messages: []Message{{Role: RoleUser, Content: "unchanged"}}, MaxTokens: 7} |
| 56 | ch, err := Stream(ctx, p, req) |
| 57 | if err == nil { |
| 58 | for chunk := range ch { |
| 59 | err = chunk.Err |
| 60 | } |
| 61 | } |
| 62 | var auth *AuthError |
| 63 | if !errors.As(err, &auth) || auth.ModelRef != "connection/model" || gate.ref != auth.ModelRef { |
| 64 | t.Fatalf("lost actual request identity: %v", err) |
| 65 | } |
| 66 | if _, err := Stream(ctx, p, req); err == nil || p.calls != 1 { |
| 67 | t.Fatal("rejection allowed automatic retry") |
| 68 | } |
| 69 | if !reflect.DeepEqual(req, p.request) { |
| 70 | t.Fatal("gate changed provider request") |
| 71 | } |
| 72 | cancel() |
| 73 | } |
| 74 | } |
| 75 |