dynagate/gateway_test.go

package main

import (
  "bytes"
  "context"
  "encoding/json"
  "fmt"
  "io"
  "net/http"
  "net/http/httptest"
  "os"
  "path/filepath"
  "reflect"
  "strings"
  "testing"
  "time"
)

func TestURLBuilding(t *testing.T) {
  tests := []struct {
    endpoint string
    reqPath  string
    expected string
  }{
    {"https://api.openai.com/v1", "/v1/chat/completions", "https://api.openai.com/v1/chat/completions"},
    {"https://api.openai.com/v1/", "/v1/chat/completions", "https://api.openai.com/v1/chat/completions"},
    {"https://api.openai.com", "/v1/chat/completions", "https://api.openai.com/v1/chat/completions"},
    {"http://localhost:8000/api/v1", "/v1/chat/completions", "http://localhost:8000/api/v1/chat/completions"},
    {"http://localhost:11434", "/v1/chat/completions", "http://localhost:11434/v1/chat/completions"},
    {"https://custom.endpoint/v1/custom", "/v1/chat/completions", "https://custom.endpoint/v1/custom/v1/chat/completions"},
  }

  for _, tc := range tests {
    t.Run(tc.endpoint+"+"+tc.reqPath, func(t *testing.T) {
      got := buildURL(tc.endpoint, tc.reqPath)
      if got != tc.expected {
        t.Errorf("buildURL(%q, %q) = %q; want %q", tc.endpoint, tc.reqPath, got, tc.expected)
      }
    })
  }
}

func TestTrialOrdering(t *testing.T) {
  configs := []ModelConfig{
    {Model: "gpt-4", Key: "key-a1", Endpoint: "ep-a"},
    {Model: "gpt-4", Key: "key-a2", Endpoint: "ep-a"},
    {Model: "gpt-3.5", Key: "key-b1", Endpoint: "ep-b"},
    {Model: "claude", Key: "key-c1", Endpoint: "ep-c"},
  }
  uniqueModels := []string{"gpt-4", "gpt-3.5", "claude"}

  // Test case 1: Requested model matches the first model "gpt-4"
  {
    trials := getTrialConfigs(configs, uniqueModels, "gpt-4")
    expected := []ModelConfig{
      {Model: "gpt-4", Key: "key-a1", Endpoint: "ep-a"},
      {Model: "gpt-4", Key: "key-a2", Endpoint: "ep-a"},
      {Model: "gpt-3.5", Key: "key-b1", Endpoint: "ep-b"},
      {Model: "claude", Key: "key-c1", Endpoint: "ep-c"},
    }
    if !reflect.DeepEqual(trials, expected) {
      t.Errorf("Trial ordering for gpt-4 mismatch.\nGot: %+v\nWant: %+v", trials, expected)
    }
  }

  // Test case 2: Requested model matches the second model "gpt-3.5"
  {
    trials := getTrialConfigs(configs, uniqueModels, "gpt-3.5")
    expected := []ModelConfig{
      {Model: "gpt-3.5", Key: "key-b1", Endpoint: "ep-b"},
      {Model: "claude", Key: "key-c1", Endpoint: "ep-c"},
      {Model: "gpt-4", Key: "key-a1", Endpoint: "ep-a"},
      {Model: "gpt-4", Key: "key-a2", Endpoint: "ep-a"},
    }
    if !reflect.DeepEqual(trials, expected) {
      t.Errorf("Trial ordering for gpt-3.5 mismatch.\nGot: %+v\nWant: %+v", trials, expected)
    }
  }

  // Test case 3: Requested model matches the last model "claude"
  {
    trials := getTrialConfigs(configs, uniqueModels, "claude")
    expected := []ModelConfig{
      {Model: "claude", Key: "key-c1", Endpoint: "ep-c"},
      {Model: "gpt-4", Key: "key-a1", Endpoint: "ep-a"},
      {Model: "gpt-4", Key: "key-a2", Endpoint: "ep-a"},
      {Model: "gpt-3.5", Key: "key-b1", Endpoint: "ep-b"},
    }
    if !reflect.DeepEqual(trials, expected) {
      t.Errorf("Trial ordering for claude mismatch.\nGot: %+v\nWant: %+v", trials, expected)
    }
  }

  // Test case 4: Requested model doesn't match any config
  {
    trials := getTrialConfigs(configs, uniqueModels, "non-existent")
    expected := []ModelConfig{
      {Model: "gpt-4", Key: "key-a1", Endpoint: "ep-a"},
      {Model: "gpt-4", Key: "key-a2", Endpoint: "ep-a"},
      {Model: "gpt-3.5", Key: "key-b1", Endpoint: "ep-b"},
      {Model: "claude", Key: "key-c1", Endpoint: "ep-c"},
    }
    if !reflect.DeepEqual(trials, expected) {
      t.Errorf("Trial ordering for non-existent model mismatch.\nGot: %+v\nWant: %+v", trials, expected)
    }
  }
}

func TestConfigManagerCSVWatcher(t *testing.T) {
  tmpDir, err := os.MkdirTemp("", "dynagate-test")
  if err != nil {
    t.Fatalf("failed to create temp dir: %v", err)
  }
  defer os.RemoveAll(tmpDir)

  csvFile := filepath.Join(tmpDir, "models.csv")

  // Write initial content
  initialContent := `model,key,endpoint
gpt-4,key-1,http://localhost:8001
gpt-3.5,key-2,http://localhost:8002
`
  if err := os.WriteFile(csvFile, []byte(initialContent), 0644); err != nil {
    t.Fatalf("failed to write csv file: %v", err)
  }

  cm := NewConfigManager(csvFile)
  configs := cm.GetConfigs()
  if len(configs) != 2 || configs[0].Model != "gpt-4" || configs[1].Model != "gpt-3.5" {
    t.Fatalf("unexpected initial configs: %+v", configs)
  }

  // Start watching in background
  ctx, cancel := context.WithCancel(context.Background())
  defer cancel()
  go cm.Watch(ctx)

  // Wait briefly, modify file, and check if reload occurs
  time.Sleep(100 * time.Millisecond)

  updatedContent := `model,key,endpoint
gpt-4,key-1,http://localhost:8001
gpt-3.5,key-2,http://localhost:8002
claude,key-3,http://localhost:8003
`
  // Make sure we sleep slightly to ensure file modtime actually changes on modern filesystems
  time.Sleep(1 * time.Second)
  if err := os.WriteFile(csvFile, []byte(updatedContent), 0644); err != nil {
    t.Fatalf("failed to update csv file: %v", err)
  }

  // Poll GetConfigs until it gets updated (timeout after 3 seconds)
  deadline := time.Now().Add(3 * time.Second)
  var updatedConfigs []ModelConfig
  for time.Now().Before(deadline) {
    updatedConfigs = cm.GetConfigs()
    if len(updatedConfigs) == 3 {
      break
    }
    time.Sleep(100 * time.Millisecond)
  }

  if len(updatedConfigs) != 3 || updatedConfigs[2].Model != "claude" {
    t.Errorf("Hot-reload failed. Got configs: %+v", updatedConfigs)
  }
}

func TestGatewayCompletionsFailover(t *testing.T) {
  // We will setup 3 mock upstream servers.
  // Server A: Always returns 500
  // Server B: Always returns 401
  // Server C: Succeeds and returns a chat completion response

  serverA := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    w.WriteHeader(http.StatusInternalServerError)
    _, _ = w.Write([]byte(`{"error": "Internal Server Error A"}`))
  }))
  defer serverA.Close()

  serverB := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    w.WriteHeader(http.StatusUnauthorized)
    _, _ = w.Write([]byte(`{"error": "Unauthorized B"}`))
  }))
  defer serverB.Close()

  serverC := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    // Read request body to verify model was modified
    var body map[string]any
    _ = json.NewDecoder(r.Body).Decode(&body)
    modelSent := body["model"].(string)

    w.Header().Set("Content-Type", "application/json")
    w.WriteHeader(http.StatusOK)
    _, _ = w.Write([]byte(fmt.Sprintf(`{"choices":[{"message":{"content":"Hello from %s"}}]}`, modelSent)))
  }))
  defer serverC.Close()

  // Configure gateway ConfigManager with endpoints pointing to these servers
  // We want to verify that when we call completions:
  // - First it tries model-1 (pointing to serverA) -> fails
  // - Then tries model-1 next key (pointing to serverB) -> fails
  // - Then retries with next model model-2 (pointing to serverC) -> succeeds!
  configs := []ModelConfig{
    {Model: "model-1", Key: "key-fail-1", Endpoint: serverA.URL},
    {Model: "model-1", Key: "key-fail-2", Endpoint: serverB.URL},
    {Model: "model-2", Key: "key-success", Endpoint: serverC.URL},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-1", "model-2"},
  }

  handler := handleChatCompletions(cm, "")

  // 1. Test non-streaming failover request starting with model-1
  reqBodyObj := map[string]any{
    "model": "model-1",
    "messages": []map[string]string{
      {"role": "user", "content": "Hi"},
    },
  }
  reqBytes, _ := json.Marshal(reqBodyObj)

  req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
  req.Header.Set("Content-Type", "application/json")
  w := httptest.NewRecorder()

  handler.ServeHTTP(w, req)

  resp := w.Result()
  defer resp.Body.Close()

  if resp.StatusCode != http.StatusOK {
    t.Fatalf("Expected status 200, got %d", resp.StatusCode)
  }

  var respObj map[string]any
  _ = json.NewDecoder(resp.Body).Decode(&respObj)

  choices, ok := respObj["choices"].([]any)
  if !ok || len(choices) == 0 {
    t.Fatalf("Invalid response structure: %+v", respObj)
  }

  content := choices[0].(map[string]any)["message"].(map[string]any)["content"].(string)
  expectedContent := "Hello from model-2"
  if content != expectedContent {
    t.Errorf("Expected content %q, got %q (failover to model-2 did not happen or model was not updated)", expectedContent, content)
  }
}

func TestGatewayCompletionsStreamingFailover(t *testing.T) {
  // We will setup 2 mock upstream servers.
  // Server A: Always returns 500
  // Server B: Streams chat completions chunks

  serverA := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    w.WriteHeader(http.StatusInternalServerError)
  }))
  defer serverA.Close()

  serverB := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    w.Header().Set("Content-Type", "text/event-stream")
    w.Header().Set("Cache-Control", "no-cache")
    w.Header().Set("Connection", "keep-alive")
    w.WriteHeader(http.StatusOK)

    flusher, _ := w.(http.Flusher)
    _, _ = w.Write([]byte("data: chunk1\n\n"))
    flusher.Flush()
    _, _ = w.Write([]byte("data: chunk2\n\n"))
    flusher.Flush()
  }))
  defer serverB.Close()

  configs := []ModelConfig{
    {Model: "model-1", Key: "key-fail", Endpoint: serverA.URL},
    {Model: "model-2", Key: "key-success", Endpoint: serverB.URL},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-1", "model-2"},
  }

  handler := handleChatCompletions(cm, "")

  reqBodyObj := map[string]any{
    "model":  "model-1",
    "stream": true,
    "messages": []map[string]string{
      {"role": "user", "content": "Hi"},
    },
  }
  reqBytes, _ := json.Marshal(reqBodyObj)

  req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
  req.Header.Set("Content-Type", "application/json")
  w := httptest.NewRecorder()

  handler.ServeHTTP(w, req)

  resp := w.Result()
  defer resp.Body.Close()

  if resp.StatusCode != http.StatusOK {
    t.Fatalf("Expected status 200, got %d", resp.StatusCode)
  }

  respBytes, _ := io.ReadAll(resp.Body)
  respStr := string(respBytes)

  if !strings.Contains(respStr, "data: chunk1") || !strings.Contains(respStr, "data: chunk2") {
    t.Errorf("Expected streaming chunks to be forwarded, got: %q", respStr)
  }
}

func TestModelsEndpoint(t *testing.T) {
  configs := []ModelConfig{
    {Model: "model-1", Key: "key-1", Endpoint: "ep-1"},
    {Model: "model-1", Key: "key-2", Endpoint: "ep-1"},
    {Model: "model-2", Key: "key-3", Endpoint: "ep-2"},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-1", "model-2"},
  }

  handler := handleModels(cm, "")

  req := httptest.NewRequest("GET", "/v1/models", nil)
  w := httptest.NewRecorder()

  handler.ServeHTTP(w, req)

  resp := w.Result()
  defer resp.Body.Close()

  if resp.StatusCode != http.StatusOK {
    t.Fatalf("Expected status 200, got %d", resp.StatusCode)
  }

  var respObj map[string]any
  _ = json.NewDecoder(resp.Body).Decode(&respObj)

  data, ok := respObj["data"].([]any)
  if !ok || len(data) != 2 {
    t.Fatalf("Expected 2 models in response data, got: %+v", respObj)
  }

  model1 := data[0].(map[string]any)["id"].(string)
  model2 := data[1].(map[string]any)["id"].(string)

  if model1 != "model-1" || model2 != "model-2" {
    t.Errorf("Unexpected models. Got: %s, %s", model1, model2)
  }
}

func TestGatewayAuthentication(t *testing.T) {
  configs := []ModelConfig{
    {Model: "model-1", Key: "key-1", Endpoint: "ep-1"},
  }
  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-1"},
  }

  // Set up handlers expecting "secret-token"
  modelsHandler := handleModels(cm, "secret-token")
  completionsHandler := handleChatCompletions(cm, "secret-token")

  // Case 1: No Authorization header
  {
    req := httptest.NewRequest("GET", "/v1/models", nil)
    w := httptest.NewRecorder()
    modelsHandler.ServeHTTP(w, req)
    if w.Result().StatusCode != http.StatusUnauthorized {
      t.Errorf("Expected 401 Unauthorized for missing token, got %d", w.Result().StatusCode)
    }
  }

  // Case 2: Incorrect Authorization header
  {
    req := httptest.NewRequest("GET", "/v1/models", nil)
    req.Header.Set("Authorization", "Bearer wrong-token")
    w := httptest.NewRecorder()
    modelsHandler.ServeHTTP(w, req)
    if w.Result().StatusCode != http.StatusUnauthorized {
      t.Errorf("Expected 401 Unauthorized for wrong token, got %d", w.Result().StatusCode)
    }
  }

  // Case 3: Correct Authorization header
  {
    req := httptest.NewRequest("GET", "/v1/models", nil)
    req.Header.Set("Authorization", "Bearer secret-token")
    w := httptest.NewRecorder()
    modelsHandler.ServeHTTP(w, req)
    if w.Result().StatusCode != http.StatusOK {
      t.Errorf("Expected 200 OK for correct token, got %d", w.Result().StatusCode)
    }
  }

  // Case 4: Completions endpoint with correct token
  {
    reqBodyObj := map[string]any{
      "model": "model-1",
    }
    reqBytes, _ := json.Marshal(reqBodyObj)
    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
    req.Header.Set("Authorization", "Bearer secret-token")
    req.Header.Set("Content-Type", "application/json")
    w := httptest.NewRecorder()
    completionsHandler.ServeHTTP(w, req)
    // It should proceed past auth and return 502/503/500 because "ep-1" is not a valid endpoint,
    // but critically, it should NOT return 401.
    status := w.Result().StatusCode
    if status == http.StatusUnauthorized {
      t.Errorf("Expected completions request to pass authentication, but got 401 Unauthorized")
    }
  }
}

func TestCSVUpdater(t *testing.T) {
  tmpDir, err := os.MkdirTemp("", "dynagate-updater-test")
  if err != nil {
    t.Fatalf("failed to create temp dir: %v", err)
  }
  defer os.RemoveAll(tmpDir)

  csvFile := filepath.Join(tmpDir, "models.csv")
  initialContent := "model,key,endpoint\ngpt-4,key-1,http://localhost:8001\n"
  if err := os.WriteFile(csvFile, []byte(initialContent), 0644); err != nil {
    t.Fatalf("failed to write csv file: %v", err)
  }

  // Create a command to append a model row to the CSV file
  command := fmt.Sprintf("echo 'gpt-3.5,key-2,http://localhost:8002' >> %s", csvFile)

  ctx, cancel := context.WithCancel(context.Background())
  defer cancel()

  // Start background updater with a fast interval (100ms)
  go startCSVUpdater(ctx, command, 100*time.Millisecond)

  // The updater runs once immediately at startup, so we expect the file to have the appended line
  // almost immediately. Let's poll to check.
  deadline := time.Now().Add(2 * time.Second)
  var contentBytes []byte
  var containsNewModel bool

  for time.Now().Before(deadline) {
    contentBytes, err = os.ReadFile(csvFile)
    if err == nil && strings.Contains(string(contentBytes), "gpt-3.5") {
      containsNewModel = true
      break
    }
    time.Sleep(50 * time.Millisecond)
  }

  if !containsNewModel {
    t.Fatalf("Expected CSV file to contain the updater-appended model, but it didn't. Content: %s", string(contentBytes))
  }
}

func TestGatewayImageGenerations(t *testing.T) {
  // We will setup a mock upstream server that handles image generation.
  server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if r.URL.Path != "/v1/images/generations" {
      w.WriteHeader(http.StatusNotFound)
      return
    }

    var body map[string]any
    _ = json.NewDecoder(r.Body).Decode(&body)
    modelSent := body["model"].(string)

    w.Header().Set("Content-Type", "application/json")
    w.WriteHeader(http.StatusOK)
    _, _ = w.Write([]byte(fmt.Sprintf(`{"created":1589478378,"data":[{"url":"http://image-url/generated-by-%s"}]}`, modelSent)))
  }))
  defer server.Close()

  configs := []ModelConfig{
    {Model: "dall-e-3", Key: "dalle-key", Endpoint: server.URL},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"dall-e-3"},
  }

  handler := handleImageGenerations(cm, "")

  reqBodyObj := map[string]any{
    "prompt": "a beautiful kitten",
    "model": "dall-e-3",
  }
  reqBytes, _ := json.Marshal(reqBodyObj)

  req := httptest.NewRequest("POST", "/v1/images/generations", bytes.NewReader(reqBytes))
  req.Header.Set("Content-Type", "application/json")
  w := httptest.NewRecorder()

  handler.ServeHTTP(w, req)

  resp := w.Result()
  defer resp.Body.Close()

  if resp.StatusCode != http.StatusOK {
    t.Fatalf("Expected status 200, got %d", resp.StatusCode)
  }

  var respObj map[string]any
  _ = json.NewDecoder(resp.Body).Decode(&respObj)

  data, ok := respObj["data"].([]any)
  if !ok || len(data) == 0 {
    t.Fatalf("Expected data field to contain image urls, got: %+v", respObj)
  }

  urlVal := data[0].(map[string]any)["url"].(string)
  expectedURL := "http://image-url/generated-by-dall-e-3"
  if urlVal != expectedURL {
    t.Errorf("Expected URL %q, got %q", expectedURL, urlVal)
  }
}

func TestGatewayNoAuthHeaderOnEmptyKey(t *testing.T) {
  server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if auth := r.Header.Get("Authorization"); auth != "" {
      t.Errorf("Expected no Authorization header, but got %q", auth)
      w.WriteHeader(http.StatusBadRequest)
      return
    }
    w.Header().Set("Content-Type", "application/json")
    w.WriteHeader(http.StatusOK)
    _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"Hello"}}]}n`))
  }))
  defer server.Close()

  configs := []ModelConfig{
    {Model: "model-no-key", Key: "", Endpoint: server.URL},
    {Model: "model-no-key-2", Key: "-", Endpoint: server.URL},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-no-key", "model-no-key-2"},
  }

  handler := handleChatCompletions(cm, "")

  // Case 1: Empty Key
  {
    reqBodyObj := map[string]any{
      "model": "model-no-key",
      "messages": []map[string]string{
        {"role": "user", "content": "Hi"},
      },
    }
    reqBytes, _ := json.Marshal(reqBodyObj)
    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
    req.Header.Set("Content-Type", "application/json")
    req.Header.Set("Authorization", "Bearer some-client-token")
    w := httptest.NewRecorder()
    handler.ServeHTTP(w, req)
    if w.Result().StatusCode != http.StatusOK {
      t.Errorf("Expected status 200 for empty key, got %d", w.Result().StatusCode)
    }
  }

  // Case 2: Dash Key
  {
    reqBodyObj := map[string]any{
      "model": "model-no-key-2",
      "messages": []map[string]string{
        {"role": "user", "content": "Hi"},
      },
    }
    reqBytes, _ := json.Marshal(reqBodyObj)
    req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
    req.Header.Set("Content-Type", "application/json")
    req.Header.Set("Authorization", "Bearer some-client-token")
    w := httptest.NewRecorder()
    handler.ServeHTTP(w, req)
    if w.Result().StatusCode != http.StatusOK {
      t.Errorf("Expected status 200 for dash key, got %d", w.Result().StatusCode)
    }
  }
}

func TestGatewayBlankKey(t *testing.T) {
  server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    auth := r.Header.Get("Authorization")
    if auth != "Bearer" {
      t.Errorf("Expected Authorization header to be exactly 'Bearer', but got %q", auth)
      w.WriteHeader(http.StatusBadRequest)
      return
    }
    w.Header().Set("Content-Type", "application/json")
    w.WriteHeader(http.StatusOK)
    _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"Hello"}}]}n`))
  }))
  defer server.Close()

  configs := []ModelConfig{
    {Model: "model-blank-key", Key: "-blank-", Endpoint: server.URL},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-blank-key"},
  }

  handler := handleChatCompletions(cm, "")

  reqBodyObj := map[string]any{
    "model": "model-blank-key",
    "messages": []map[string]string{
      {"role": "user", "content": "Hi"},
    },
  }
  reqBytes, _ := json.Marshal(reqBodyObj)

  req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
  req.Header.Set("Content-Type", "application/json")
  req.Header.Set("Authorization", "Bearer some-client-token")
  w := httptest.NewRecorder()

  handler.ServeHTTP(w, req)

  if w.Result().StatusCode != http.StatusOK {
    t.Errorf("Expected status 200 for -blank- key, got %d", w.Result().StatusCode)
  }
}

func TestGatewayExtraHeaders(t *testing.T) {
  server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if h := r.Header.Get("X-My-Custom-Header"); h != "HeaderValue" {
      t.Errorf("Expected X-My-Custom-Header to be 'HeaderValue', got %q", h)
      w.WriteHeader(http.StatusBadRequest)
      return
    }
    if h := r.Header.Get("X-Another-Header"); h != "123" {
      t.Errorf("Expected X-Another-Header to be '123', got %q", h)
      w.WriteHeader(http.StatusBadRequest)
      return
    }
    w.Header().Set("Content-Type", "application/json")
    w.WriteHeader(http.StatusOK)
    _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"Hello"}}]}n`))
  }))
  defer server.Close()

  configs := []ModelConfig{
    {Model: "model-extra-headers", Key: "my-key", Endpoint: server.URL, Extra: `{"X-My-Custom-Header":"HeaderValue","X-Another-Header":123}`},
  }

  cm := &ConfigManager{
    configs:      configs,
    uniqueModels: []string{"model-extra-headers"},
  }

  handler := handleChatCompletions(cm, "")

  reqBodyObj := map[string]any{
    "model": "model-extra-headers",
    "messages": []map[string]string{
      {"role": "user", "content": "Hi"},
    },
  }
  reqBytes, _ := json.Marshal(reqBodyObj)

  req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewReader(reqBytes))
  req.Header.Set("Content-Type", "application/json")
  w := httptest.NewRecorder()

  handler.ServeHTTP(w, req)

  if w.Result().StatusCode != http.StatusOK {
    t.Errorf("Expected status 200, got %d", w.Result().StatusCode)
  }
}