CrossRelay/internal/chatcmpl/client_test.go

package chatcmpl

import (
	"context"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestCreateNonStream(t *testing.T) {
	var gotAuth, gotBody string
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotAuth = r.Header.Get("Authorization")
		data, _ := io.ReadAll(r.Body)
		gotBody = string(data)
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"id":"cmpl-1","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}`))
	}))
	defer upstream.Close()

	client := NewClient(upstream.URL, "sk-up")
	resp, err := client.Create(context.Background(), &Request{Model: "ds-chat", Messages: []Message{{Role: "user", Content: "hi"}}})
	if err != nil {
		t.Fatalf("Create: %v", err)
	}
	if gotAuth != "Bearer sk-up" {
		t.Errorf("authorization = %q", gotAuth)
	}
	if !strings.Contains(gotBody, `"model":"ds-chat"`) {
		t.Errorf("request body = %s", gotBody)
	}
	if len(resp.Choices) != 1 || resp.Choices[0].Message.Content != "hi" {
		t.Fatalf("response = %+v", resp)
	}
	if resp.Usage == nil || resp.Usage.PromptTokens != 5 {
		t.Errorf("usage = %+v", resp.Usage)
	}
}

func TestCreateUpstreamError(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusUnauthorized)
		_, _ = w.Write([]byte(`{"error":{"message":"bad key","type":"invalid_api_key"}}`))
	}))
	defer upstream.Close()

	client := NewClient(upstream.URL, "sk-bad")
	_, err := client.Create(context.Background(), &Request{Model: "m"})
	var ue *UpstreamError
	if !errors.As(err, &ue) {
		t.Fatalf("error = %v, want *UpstreamError", err)
	}
	if ue.StatusCode != http.StatusUnauthorized || ue.Message != "bad key" {
		t.Errorf("upstream error = %+v", ue)
	}
}

func TestCreateStreamRejectsNon2xx(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusTooManyRequests)
		_, _ = w.Write([]byte(`{"error":{"message":"slow down"}}`))
	}))
	defer upstream.Close()

	client := NewClient(upstream.URL, "sk")
	resp, err := client.CreateStream(context.Background(), &Request{Model: "m", Stream: true})
	if err == nil {
		resp.Body.Close()
		t.Fatal("CreateStream succeeded on 429")
	}
}

func TestStreamNext(t *testing.T) {
	sse := strings.Join([]string{
		`data: {"id":"1","choices":[{"index":0,"delta":{"content":"Hel"}}]}`,
		"",
		`: keep-alive comment`,
		`data: {"id":"1","choices":[{"index":0,"delta":{"content":"lo"}}]}`,
		"",
		`data: {"id":"1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":1}}`,
		"",
		"data: [DONE]",
		"",
	}, "\n")
	stream := NewStream(strings.NewReader(sse))

	first, err := stream.Next()
	if err != nil {
		t.Fatalf("Next: %v", err)
	}
	if first.Choices[0].Delta.Content == nil || *first.Choices[0].Delta.Content != "Hel" {
		t.Fatalf("first chunk = %+v", first)
	}
	second, err := stream.Next()
	if err != nil {
		t.Fatalf("Next: %v", err)
	}
	if *second.Choices[0].Delta.Content != "lo" {
		t.Fatalf("second chunk = %+v", second)
	}
	third, err := stream.Next()
	if err != nil {
		t.Fatalf("Next: %v", err)
	}
	if third.Choices[0].FinishReason == nil || *third.Choices[0].FinishReason != "stop" {
		t.Fatalf("third chunk = %+v", third)
	}
	if third.Usage == nil || third.Usage.CompletionTokens != 1 {
		t.Fatalf("usage chunk = %+v", third)
	}
	if _, err := stream.Next(); err != io.EOF {
		t.Fatalf("final Next err = %v, want io.EOF", err)
	}
}