gemini_test.go

Overview

This file provides unit and integration tests for the Gemini model implementation within the gemini package. The Gemini model is an LLM (Large Language Model) client interface designed to interact with Google’s Gemini language models using the genai SDK. The tests focus on validating the model's content generation capabilities, streaming response handling, and proper inclusion of HTTP tracking headers during API calls.

The file also contains supporting utility code for test client configuration, response reading/parsing from streaming LLM output, and an HTTP interceptor to verify request headers.

Detailed Breakdown

Test Functions

TestModel_Generate(t *testing.T)

• Purpose:
  Tests the synchronous generation of LLM responses using the Gemini model.

• Parameters:

  • t *testing.T: Standard Go testing object.

• Functionality:
  The test defines a table of test cases with input requests (model.LLMRequest) and expected responses (model.LLMResponse). For each test case, it:

  1. Creates a new Gemini model instance with a test HTTP client configuration that supports record/replay (httprr).

  2. Calls GenerateContent() on the model to synchronously generate content.

  3. Compares the generated response with the expected response using cmp.Diff, ignoring average log probabilities.

  4. Validates error presence matches expectations.

• Important Notes:
  The test uses HTTP recording files stored under testdata to mock API interactions for consistent testing.

• Example Usage:

  req := &model.LLMRequest{
      Contents: genai.Text("What is the capital of France? One word."),
      Config: &genai.GenerateContentConfig{
          Temperature: new(float32),
      },
  }
  // Setup and call GenerateContent, then validate output.
  

TestModel_GenerateStream(t *testing.T)

• Purpose:
  Tests streaming generation where LLM content is returned incrementally.

• Parameters:

  • t *testing.T: Standard Go testing object.

• Functionality:
  Similar setup to TestModel_Generate, but calls GenerateContent() with streaming enabled (true). It then uses readResponse() to concatenate partial streaming responses and compares the accumulated text to expected results.

• Checks:

  • Correct concatenation of partial and final texts.

  • Error handling consistency.

TestModel_TrackingHeaders(t *testing.T)

• Purpose:
  Verifies that HTTP requests made by the Gemini model include necessary tracking headers for telemetry.

• Parameters:

  • t *testing.T: Standard Go testing object.

• Functionality:

  • Sets up a custom headerInterceptor which inspects outbound HTTP requests.

  • Checks for presence and content of "User-Agent" and "x-goog-api-client" headers.

  • Uses a fake API key if not recording live HTTP traffic.

  • Triggers a dummy request via the Gemini model.

  • Fails test if headers are missing or malformed.

Helper Functions and Types

newGeminiTestClientConfig(t *testing.T, rrfile string) *genai.ClientConfig

• Purpose:
  Constructs a genai.ClientConfig configured for HTTP record/replay testing.

• Parameters:

  • t *testing.T: Testing handle for error reporting.

  • rrfile string: Path to HTTP record/replay file.

• Returns:

  • *genai.ClientConfig: Configured client.

• Implementation Detail:
  Uses testutil.NewGeminiTransport to create a transport that replays or records HTTP requests. Sets a fake API key if not recording.

type TextResponse

• Description:
  Container struct holding concatenated text from streaming LLM responses, split into partial (streaming) and final (complete) segments.

• Fields:

  • PartialText string: Concatenation of all streamed partial response texts.

  • FinalText string: Concatenation of all final response texts.

func readResponse(s iter.Seq2[*model.LLMResponse, error]) (TextResponse, error)

• Purpose:
  Aggregates streamed LLM responses from a sequence into a TextResponse, separating partial and final content.

• Parameters:

  • s iter.Seq2[*model.LLMResponse, error]: A sequence that yields LLM responses and possible errors.

• Returns:

  • TextResponse: Aggregated partial and final texts.

  • error: Any error encountered during reading.

• Implementation Detail:
  Iterates over the sequence, concatenates the first text part of each response to either partial or final strings depending on the Partial flag. Returns early if an error or empty response is encountered.

type headerInterceptor

• Description:
  An HTTP RoundTripper wrapper that executes a user-provided check function on each request before delegating to the base RoundTripper.

• Fields:

  • base http.RoundTripper: The underlying RoundTripper to delegate to.

  • check func(*http.Request): Function to invoke on each request for inspection or modification.

• Methods:

  • RoundTrip(req *http.Request) (*http.Response, error):
    Executes check on the request and then forwards it to base.RoundTrip. Uses http.DefaultTransport if base is nil.

Implementation Details and Algorithms

• Test Coverage Using HTTP Record/Replay:
  Tests rely on HTTP record/replay files (.httprr) stored in testdata to simulate API responses. This allows for deterministic and offline-capable tests.

• Response Streaming Handling:
  Streaming responses are handled by iterating over a sequence of partial LLM responses, distinguishing partial vs final responses. The readResponse() utility consolidates these into full text strings, enabling validation of streamed output.

• Header Interception:
  The headerInterceptor wraps HTTP transport to validate that telemetry/tracking headers are included in outgoing requests without modifying the underlying transport logic.

• Use of Go Context:
  Model generation calls use t.Context() to obtain contextual cancellation and timeout signals, ensuring tests respect the testing lifecycle.

Interactions with Other Parts of the System

• Model Implementation (NewModel):
  The tests instantiate Gemini model clients via NewModel(). This Factory function is assumed to be implemented elsewhere in the gemini package and provides configured LLM client instances.

• Model Request and Response Types:
  Uses model.LLMRequest and model.LLMResponse types from the model package to structure input requests and expected responses.

• genai SDK Integration:
  The file heavily utilizes the genai package for request/response content, configuration, and client setup, integrating Gemini models into the Google AI ecosystem.

• HTTP Record/Replay Support:
  Uses httprr and testutil packages to support deterministic HTTP interaction recording and playback during tests.

Usage Example (Simplified)

// Create Gemini model client with test config
cfg := newGeminiTestClientConfig(t, "testdata/TestModel_Generate.ok.httprr")
model, err := NewModel(t.Context(), "gemini-2.0-flash", cfg)
if err != nil {
    t.Fatal(err)
}

// Prepare LLM request
req := &model.LLMRequest{
    Contents: genai.Text("What is the capital of France?"),
    Config: &genai.GenerateContentConfig{Temperature: new(float32)},
}

// Generate response synchronously
for resp, err := range model.GenerateContent(t.Context(), req, false) {
    if err != nil {
        t.Fatal(err)
    }
    fmt.Println(resp.Content.Parts[0].Text) // Expected "Paris\n"
}

Diagram: Flowchart of Test Functions and Utility Relationships

flowchart TD
subgraph Tests
TG[TestModel_Generate]
TS[TestModel_GenerateStream]
TH[TestModel_TrackingHeaders]
end
subgraph Utilities
NG[newGeminiTestClientConfig]
RR[readResponse]
HI[headerInterceptor]
end
TG --> NG
TG -->|calls| NewModel
TS --> NG
TS -->|calls| NewModel
TS --> RR
TH --> NG
TH --> HI
HI -->|wraps| HTTPTransport
NewModel[NewModel Factory]
HTTPTransport[http.RoundTripper]
style Tests fill:none,stroke:#333,stroke-width:1px
style Utilities fill:none,stroke:#333,stroke-width:1px

• This diagram illustrates how the test functions depend on utility functions and types.

• TestModel_Generate and TestModel_GenerateStream both rely on newGeminiTestClientConfig and the NewModel factory to create test clients.

• TestModel_GenerateStream uses readResponse to aggregate streaming responses.

• TestModel_TrackingHeaders employs headerInterceptor to inspect HTTP requests.

• The headerInterceptor wraps an underlying HTTP transport to intercept and check request headers.

This documentation describes the testing logic and support mechanisms for ensuring the Gemini model's API client behaves correctly in generating content, streaming responses, and including telemetry headers during HTTP requests. It is a critical part of verifying the integration and correctness of the Gemini LLM client within the system.