Vectrum
Sign In
186 lines · 5.9 KB · Go
package llm

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

const okResponse = `{"id":"msg_1","type":"message","role":"assistant","model":"claude-opus-5",
"content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","usage":{"input_tokens":3,"output_tokens":2}}`

func testClient(t *testing.T, h http.HandlerFunc) *Client {
	t.Helper()
	srv := httptest.NewServer(h)
	t.Cleanup(srv.Close)
	c := New(WithBaseURL(srv.URL), WithAPIKey("test-key"))
	c.RetryDelay = time.Millisecond // keep the retry tests fast
	return c
}

func TestCreateMessageSendsExpectedHeadersAndBody(t *testing.T) {
	var gotHeaders http.Header
	var gotBody map[string]any
	c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
		gotHeaders = r.Header.Clone()
		body, _ := io.ReadAll(r.Body)
		json.Unmarshal(body, &gotBody)
		if r.URL.Path != "/v1/messages" {
			t.Errorf("path = %q, want /v1/messages", r.URL.Path)
		}
		io.WriteString(w, okResponse)
	})

	resp, err := c.CreateMessage(context.Background(), Request{
		Messages: []Message{{Role: RoleUser, Content: []ContentBlock{TextBlock("hello")}}},
	})
	if err != nil {
		t.Fatalf("CreateMessage: %v", err)
	}
	if resp.Text() != "hi" {
		t.Errorf("Text() = %q, want hi", resp.Text())
	}
	if gotHeaders.Get("x-api-key") != "test-key" {
		t.Errorf("x-api-key = %q", gotHeaders.Get("x-api-key"))
	}
	if gotHeaders.Get("anthropic-version") != apiVersion {
		t.Errorf("anthropic-version = %q, want %s", gotHeaders.Get("anthropic-version"), apiVersion)
	}
	if gotBody["model"] != DefaultModel {
		t.Errorf("model = %v, want the client default %s", gotBody["model"], DefaultModel)
	}
	if gotBody["max_tokens"].(float64) != DefaultMaxTokens {
		t.Errorf("max_tokens = %v, want the client default", gotBody["max_tokens"])
	}
}

func TestOAuthTokenUsesBearerAndBetaHeader(t *testing.T) {
	var got http.Header
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got = r.Header.Clone()
		io.WriteString(w, okResponse)
	}))
	defer srv.Close()

	c := New(WithBaseURL(srv.URL), WithAuthToken("oauth-tok"))
	if _, err := c.CreateMessage(context.Background(), Request{}); err != nil {
		t.Fatalf("CreateMessage: %v", err)
	}
	if got.Get("authorization") != "Bearer oauth-tok" {
		t.Errorf("authorization = %q, want Bearer (OAuth tokens do not go on x-api-key)", got.Get("authorization"))
	}
	if got.Get("x-api-key") != "" {
		t.Errorf("x-api-key should be unset when using an OAuth token, got %q", got.Get("x-api-key"))
	}
	if got.Get("anthropic-beta") != oauthBeta {
		t.Errorf("anthropic-beta = %q, want %s", got.Get("anthropic-beta"), oauthBeta)
	}
}

func TestNoCredentialsIsAClearError(t *testing.T) {
	c := New(WithBaseURL("http://127.0.0.1:1"))
	c.apiKey, c.authToken = "", ""
	_, err := c.CreateMessage(context.Background(), Request{})
	if err == nil || !strings.Contains(err.Error(), "ANTHROPIC_API_KEY") {
		t.Errorf("error = %v, want it to name the env var", err)
	}
	if c.HasCredentials() {
		t.Error("HasCredentials should be false")
	}
}

func TestRetriesTransientStatusesAndNotClientErrors(t *testing.T) {
	tests := []struct {
		name        string
		status      int
		wantCalls   int32
		wantSucceed bool
	}{
		{name: "429 is retried", status: http.StatusTooManyRequests, wantCalls: 3},
		{name: "500 is retried", status: http.StatusInternalServerError, wantCalls: 3},
		{name: "400 is not retried", status: http.StatusBadRequest, wantCalls: 1},
		{name: "401 is not retried", status: http.StatusUnauthorized, wantCalls: 1},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			var calls int32
			c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
				atomic.AddInt32(&calls, 1)
				w.WriteHeader(tc.status)
				io.WriteString(w, `{"type":"error","error":{"type":"boom","message":"nope"}}`)
			})
			c.MaxRetries = 2
			c.HTTP.Timeout = 5 * time.Second

			_, err := c.CreateMessage(context.Background(), Request{})
			if err == nil {
				t.Fatal("want an error")
			}
			if got := atomic.LoadInt32(&calls); got != tc.wantCalls {
				t.Errorf("server saw %d calls, want %d", got, tc.wantCalls)
			}
			apiErr, ok := err.(*APIError)
			if !ok {
				t.Fatalf("error type = %T, want *APIError", err)
			}
			if apiErr.StatusCode != tc.status || apiErr.Message == "" {
				t.Errorf("APIError = %+v, want the status and parsed message", apiErr)
			}
		})
	}
}

func TestRetrySucceedsOnSecondAttempt(t *testing.T) {
	var calls int32
	c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
		if atomic.AddInt32(&calls, 1) == 1 {
			w.WriteHeader(http.StatusServiceUnavailable)
			return
		}
		io.WriteString(w, okResponse)
	})
	resp, err := c.CreateMessage(context.Background(), Request{})
	if err != nil {
		t.Fatalf("CreateMessage: %v", err)
	}
	if resp.Text() != "hi" || atomic.LoadInt32(&calls) != 2 {
		t.Errorf("got %q after %d calls, want hi after 2", resp.Text(), calls)
	}
}

func TestContextCancellationStopsRetrying(t *testing.T) {
	c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
	})
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err := c.CreateMessage(ctx, Request{}); err == nil {
		t.Fatal("want an error from the cancelled context")
	}
}

func TestMalformedResponseBodyIsReported(t *testing.T) {
	c := testClient(t, func(w http.ResponseWriter, r *http.Request) {
		io.WriteString(w, `{"content":`)
	})
	_, err := c.CreateMessage(context.Background(), Request{})
	if err == nil || !strings.Contains(err.Error(), "decoding response") {
		t.Errorf("error = %v, want a decode error", err)
	}
}

func TestRetryAfterHintIsParsed(t *testing.T) {
	got, ok := retryAfter("rate limited (retry-after 7s)")
	if !ok || got != 7*time.Second {
		t.Errorf("retryAfter = (%v, %v), want (7s, true)", got, ok)
	}
	if _, ok := retryAfter("rate limited"); ok {
		t.Error("want no hint when the header was absent")
	}
}