Function TestIntegrationWithMetrics

enterprise/aibridged_integration_test.go:325–427

TestIntegrationWithMetrics validates that Prometheus metrics are correctly incremented when requests are processed through aibridged.

(t *testing.T)

// TestIntegrationWithMetrics validates that Prometheus metrics are correctly incremented
// when requests are processed through aibridged.
func TestIntegrationWithMetrics(t *testing.T) {
	t.Parallel()

	ctx := testutil.Context(t, testutil.WaitLong)

	// Create prometheus registry and metrics.
	registry := prometheus.NewRegistry()
	metrics := aibridge.NewMetrics(registry)

	// Set up mock OpenAI server.
	mockOpenAI := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(`{
  "id": "chatcmpl-test",
  "object": "chat.completion",
  "created": 1753343279,
  "model": "gpt-4.1",
  "choices": [
    {
      "index": 0,
      "message": {
        "role": "assistant",
        "content": "test response"
      },
      "finish_reason": "stop"
    }
  ],
  "usage": {
    "prompt_tokens": 10,
    "completion_tokens": 5,
    "total_tokens": 15
  }
}`))
	}))
	t.Cleanup(mockOpenAI.Close)

	// Database and coderd setup.
	db, ps := dbtestutil.NewDB(t)
	client, _, api, firstUser := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
		Options: &coderdtest.Options{
			Database: db,
			Pubsub:   ps,
		},
	})

	userClient, _ := coderdtest.CreateAnotherUser(t, client, firstUser.OrganizationID)

	// Create an API token for the user.
	apiKey, err := userClient.CreateToken(ctx, "me", codersdk.CreateTokenRequest{
		TokenName: fmt.Sprintf("test-key-%d", time.Now().UnixNano()),
		Lifetime:  time.Hour,
		Scope:     codersdk.APIKeyScopeCoderAll,
	})
	require.NoError(t, err)

	// Create aibridge client.
	aiBridgeClient, err := api.AGPL.CreateInMemoryAIBridgeServer(ctx)