stream_aggregator_test.go

Overview

This file contains unit tests for validating the streaming aggregation behavior of a language model interface. It simulates multiple calls to a streaming generation method and verifies that intermediate and accumulated responses are correctly aggregated and flagged as partial or complete. The tests cover scenarios including multiple streaming calls, intermediate resets caused by empty responses or multimedia content, and handling of purely audio/video streams without aggregation.

The core focus is on verifying how partial streaming responses are combined and when the aggregation state resets, ensuring that the aggregation logic behaves as expected under various simulated streaming conditions.

Types and Structs

`streamAggregatorTest`

This struct defines a test case for the stream aggregator tests.

Field	Type	Description
name	string	Descriptive name of the test case.
`initialResponses`	[]*genai.Content	Initial set of responses provided by the mock model to simulate streaming responses.
`numberOfStreamCalls`	`int`	Number of separate streaming calls to simulate in the test.
`streamResponsesCount`	`int`	Number of responses expected per streaming call.
`want`	[]*genai.Content	Expected aggregated stream responses after processing.
`wantPartial`	[]bool	Expected partial flag for each response indicating whether it is an intermediate partial response or a complete one.

Functions

`TestStreamAggregator`

func TestStreamAggregator(t *testing.T)

Description

This is the main test function that runs multiple subtests to validate the streaming aggregation behavior. Each subtest uses a streamAggregatorTest struct instance to provide inputs and expected outputs.

Parameters

t *testing.T: The testing context object used to report test failures and run subtests.

Behavior

Sets up multiple test cases with different streaming response sequences.
For each test case:
- A mock model is initialized using testutil.MockModel with pre-defined streaming responses.
- The test simulates multiple streaming calls (numberOfStreamCalls).
- For each generated stream response, it verifies:
  - The content matches the expected aggregated content.
  - The partial flag matches the expected partial status.
- Checks that the number of received responses matches expectations.

Usage Example

func TestStreamAggregator(t *testing.T) {
    // Runs predefined test cases such as multiple streams with resets and audio streams.
}

Implementation Details and Algorithms

Mock Streaming Model: The test uses testutil.MockModel which simulates a model's streaming generation by returning sequential responses from initialResponses. The mock returns streamResponsesCount items per stream call, allowing multiple calls to simulate sequences of streaming interactions.
Aggregation Verification: Each test case defines expected outputs in want and whether each output is partial (wantPartial). The test verifies that the model's streaming aggregation correctly concatenates preceding responses to form aggregated outputs and sets the partial flag properly.
Reset Conditions: Certain inputs cause the streaming aggregation to reset:
- nil content entries signify an empty context that resets aggregation.
- Content containing audio/video metadata (e.g., VideoMetadata) triggers reset.
These cases test that the aggregator restarts accumulation after such resets, ensuring no cross-contamination of aggregation across resets.
Audio/Video Special Handling: Streams containing only audio/video parts do not produce aggregated responses, and the partial flag is always false. This behavior is tested to confirm that aggregation logic respects media types.

Interaction with Other Components

genai Package: Defines the Content type and methods like NewContentFromText and NewContentFromParts used to create simulated responses.
testutil.MockModel: A mock implementation of a language model that simulates streaming generation, providing controlled, repeatable outputs for testing aggregation logic.
model.LLMRequest: Represents the request to the language model during streaming calls, although specifics are not exercised in these tests.
cmp Package: Used for deep comparison of expected vs. actual content responses to validate correctness.

This file is part of internal testing for streaming aggregation functionality likely implemented in the core model interaction or streaming layer, which ties to the broader LLM Integration and Agents topic.

Mermaid Diagram: Function Flow and Test Structure

flowchart TD
A[TestStreamAggregator] --> B{For each Test Case}
B --> C[Initialize MockModel with initialResponses]
C --> D[Loop over numberOfStreamCalls]
D --> E[Call GenerateStream -> Receive stream responses]
E --> F[For each response]
F --> G{Check for errors}
G -->|Error| H[Test fails]
G -->|No Error| I[Compare response content with want]
I --> J[Compare partial flag with wantPartial]
J --> K{All responses checked?}
K -->|No| E
K -->|Yes| L{All stream calls done?}
L -->|No| D
L -->|Yes| M[Verify total responses count]
M --> N[Test passes or fails]

Summary of Key Points

Tests simulate multiple streaming calls with predefined responses to verify aggregation logic.
Aggregation accumulates text responses, resetting on nil or multimedia content.
Partial flags indicate whether a response is intermediate (true) or final (false).
Audio/video-only streams do not trigger aggregation or partial responses.
The mock model and genai.Content are central to simulating streaming responses.
The test validates both content correctness and the partial status per response.