go-openai icon indicating copy to clipboard operation
go-openai copied to clipboard

Add dependency injection to ChatCompletionStream for improved testability

Open rjcorwin opened this issue 4 months ago • 0 comments

Describe the change This PR refactors the ChatCompletionStream to use dependency injection by introducing a ChatStreamReader interface. This allows for injecting custom stream readers, primarily for testing purposes, making the streaming functionality more testable and maintainable.

Describe your solution The changes include:

  • Added a ChatStreamReader interface that defines the contract for reading chat completion streams
  • Refactored ChatCompletionStream to use composition with a ChatStreamReader instead of embedding streamReader
  • Added NewChatCompletionStream() constructor function to enable dependency injection
  • Implemented explicit delegation methods (Recv(), Close(), Header(), GetRateLimitHeaders()) on ChatCompletionStream
  • Added interface compliance check via var _ ChatStreamReader = (*streamReader[ChatCompletionStreamResponse])(nil)

This approach maintains backward compatibility while enabling easier mocking and testing of streaming functionality.

Tests Added comprehensive tests demonstrating the new functionality:

  • TestChatCompletionStream_MockInjection: Tests basic mock injection with the new constructor
  • mock_streaming_demo_test.go: A complete demonstration file showing how to create mock clients and stream readers for testing, including:
    • MockOpenAIStreamClient: Full mock client implementation
    • mockStreamReader: Custom stream reader for controlled test responses
    • TestMockOpenAIStreamClient_Demo: Demonstrates assembling multiple stream chunks
    • TestMockOpenAIStreamClient_ErrorHandling: Shows error handling patterns

Additional context This refactoring improves the testability of code that depends on go-openai streaming without introducing breaking changes. The existing public API remains unchanged, but now supports dependency injection for testing scenarios. The new demo test file serves as documentation for users who want to mock streaming responses in their own tests.

rjcorwin avatar Jun 10 '25 08:06 rjcorwin