diff --git a/tests/unit_tests/test_artifacts.py b/tests/unit_tests/test_artifacts.py new file mode 100644 index 00000000..17cebeca --- /dev/null +++ b/tests/unit_tests/test_artifacts.py @@ -0,0 +1,209 @@ +"""High-value tests for artifact classes - focusing on real-world usage patterns.""" + +from typing import Any + +import pytest +from pydantic import ValidationError + +from celeste.artifacts import Artifact, AudioArtifact, ImageArtifact, VideoArtifact +from celeste.mime_types import AudioMimeType, ImageMimeType, VideoMimeType + + +class TestArtifact: + """Test base Artifact class behavior.""" + + @pytest.mark.smoke + @pytest.mark.parametrize( + "storage_combo,expected", + [ + # Single storage types + ({"url": "https://example.com/file"}, True), + ({"data": b"data"}, True), + ({"path": "/path/to/file"}, True), + # Empty artifact + ({}, False), + # Multiple storage types + ({"url": "https://example.com", "data": b"data"}, True), + ({"url": "https://example.com", "data": b"data", "path": "/path"}, True), + # Edge case: all None explicitly + ({"url": None, "data": None, "path": None}, False), + # Edge case: empty string path (common mistake) + ({"path": ""}, False), + ], + ids=[ + "single_url", + "single_data", + "single_path", + "empty_artifact", + "url_and_data", + "all_storage_types", + "all_none", + "empty_string_path", + ], + ) + def test_has_content_with_storage_combinations( + self, storage_combo: dict[str, Any], expected: bool + ) -> None: + """Test has_content correctly identifies content across all storage combinations.""" + artifact = Artifact(**storage_combo) + assert artifact.has_content == expected + + def test_artifact_with_multiple_storage_types_preserves_values(self) -> None: + """Artifact can have multiple storage types simultaneously (common in caching scenarios).""" + artifact = Artifact( + url="https://example.com/file.png", + data=b"cached data", + path="/cache/file.png", + ) + assert artifact.has_content is True + assert artifact.url == "https://example.com/file.png" + assert artifact.data == b"cached data" + assert artifact.path == "/cache/file.png" + + @pytest.mark.parametrize( + "kwargs", + [ + {"url": " "}, + {"path": " "}, + {"url": " ", "path": " "}, + ], + ids=["whitespace_url", "whitespace_path", "whitespace_both"], + ) + def test_has_content_with_whitespace_only_strings( + self, kwargs: dict[str, Any] + ) -> None: + """Whitespace-only strings should be treated as empty content.""" + assert Artifact(**kwargs).has_content is False + + def test_has_content_with_empty_string_url(self) -> None: + """Empty string URLs should be treated as no content.""" + assert Artifact(url="").has_content is False + + def test_has_content_with_empty_bytes(self) -> None: + """Empty bytes should be treated as no content.""" + assert Artifact(data=b"").has_content is False + + def test_artifact_metadata_default_behavior(self) -> None: + """Artifact metadata defaults to empty dict.""" + artifact = Artifact(url="https://example.com/file.png") + assert artifact.metadata == {} + assert isinstance(artifact.metadata, dict) + + def test_artifact_metadata_setting_values(self) -> None: + """Artifact metadata can store custom key-value pairs.""" + artifact = Artifact( + url="https://example.com/file.png", + metadata={"width": 1920, "height": 1080, "format": "png"}, + ) + assert artifact.metadata["width"] == 1920 + assert artifact.metadata["height"] == 1080 + assert artifact.metadata["format"] == "png" + + def test_artifact_metadata_is_mutable(self) -> None: + """Artifact metadata dict can be modified after creation.""" + artifact = Artifact(url="https://example.com/file.png") + artifact.metadata["custom"] = "value" + assert artifact.metadata["custom"] == "value" + + def test_artifact_with_none_mime_type(self) -> None: + """Artifact can have None mime_type.""" + artifact = Artifact(url="https://example.com/file.png", mime_type=None) + assert artifact.mime_type is None + assert artifact.has_content is True + + +class TestImageArtifact: + """Test ImageArtifact specific behavior.""" + + def test_image_artifact_accepts_image_mime_type(self) -> None: + """ImageArtifact accepts valid image MIME types.""" + artifact = ImageArtifact( + url="https://example.com/image.png", mime_type=ImageMimeType.PNG + ) + assert artifact.mime_type == ImageMimeType.PNG + + def test_image_artifact_preserves_webp_mime_type(self) -> None: + """ImageArtifact preserves WEBP MIME type.""" + artifact = ImageArtifact(data=b"webp data", mime_type=ImageMimeType.WEBP) + assert artifact.mime_type == ImageMimeType.WEBP + + def test_image_artifact_accepts_none_mime_type(self) -> None: + """ImageArtifact can have None mime_type.""" + artifact = ImageArtifact(url="https://example.com/image.png", mime_type=None) + assert artifact.mime_type is None + + @pytest.mark.parametrize( + "invalid_mime_type", + [VideoMimeType.MP4, AudioMimeType.MP3], + ids=["video_mime_type", "audio_mime_type"], + ) + def test_image_artifact_rejects_invalid_mime_type( + self, invalid_mime_type: VideoMimeType | AudioMimeType + ) -> None: + """ImageArtifact rejects non-image MIME types.""" + with pytest.raises(ValidationError, match=r".*mime_type.*"): + ImageArtifact( + url="https://example.com/file.png", + mime_type=invalid_mime_type, # type: ignore[arg-type] + ) + + +class TestVideoArtifact: + """Test VideoArtifact specific behavior.""" + + def test_video_artifact_accepts_video_mime_type(self) -> None: + """VideoArtifact accepts valid video MIME types.""" + artifact = VideoArtifact(path="/videos/sample.mp4", mime_type=VideoMimeType.MP4) + assert artifact.mime_type == VideoMimeType.MP4 + + def test_video_artifact_accepts_none_mime_type(self) -> None: + """VideoArtifact can have None mime_type.""" + artifact = VideoArtifact(path="/videos/sample.mp4", mime_type=None) + assert artifact.mime_type is None + + @pytest.mark.parametrize( + "invalid_mime_type", + [ImageMimeType.PNG, AudioMimeType.MP3], + ids=["image_mime_type", "audio_mime_type"], + ) + def test_video_artifact_rejects_invalid_mime_type( + self, invalid_mime_type: ImageMimeType | AudioMimeType + ) -> None: + """VideoArtifact rejects non-video MIME types.""" + with pytest.raises(ValidationError, match=r".*mime_type.*"): + VideoArtifact(path="/videos/sample.mp4", mime_type=invalid_mime_type) # type: ignore[arg-type] + + +class TestAudioArtifact: + """Test AudioArtifact specific behavior.""" + + @pytest.mark.parametrize("mime_type", [AudioMimeType.MP3, AudioMimeType.WAV]) + def test_audio_artifact_supports_common_formats( + self, mime_type: AudioMimeType + ) -> None: + """AudioArtifact supports common audio formats.""" + artifact = AudioArtifact( + url="https://example.com/audio", + mime_type=mime_type, + ) + assert artifact.mime_type == mime_type + + def test_audio_artifact_accepts_none_mime_type(self) -> None: + """AudioArtifact can have None mime_type.""" + artifact = AudioArtifact(url="https://example.com/audio.mp3", mime_type=None) + assert artifact.mime_type is None + + @pytest.mark.parametrize( + "invalid_mime_type", + [ImageMimeType.PNG, VideoMimeType.MP4], + ids=["image_mime_type", "video_mime_type"], + ) + def test_audio_artifact_rejects_invalid_mime_type( + self, invalid_mime_type: ImageMimeType | VideoMimeType + ) -> None: + """AudioArtifact rejects non-audio MIME types.""" + with pytest.raises(ValidationError, match=r".*mime_type.*"): + AudioArtifact( + url="https://example.com/audio.mp3", + mime_type=invalid_mime_type, # type: ignore[arg-type] + ) diff --git a/tests/unit_tests/test_client.py b/tests/unit_tests/test_client.py new file mode 100644 index 00000000..37c32a03 --- /dev/null +++ b/tests/unit_tests/test_client.py @@ -0,0 +1,511 @@ +"""High-value tests for Client - focusing on critical validation and framework behavior.""" + +from collections.abc import Generator +from enum import StrEnum +from typing import Any, Unpack + +import pytest +from pydantic import SecretStr, ValidationError + +from celeste.client import Client, _clients, get_client_class, register_client +from celeste.core import Capability, Provider +from celeste.io import Input, Output, Usage +from celeste.models import Model +from celeste.parameters import ParameterMapper, Parameters + + +class ParamEnum(StrEnum): + """Test parameter enum for unit tests.""" + + TEST_PARAM = "test_param" + FIRST_PARAM = "first_param" + SECOND_PARAM = "second_param" + + +class _TestInput(Input): + """Test input with prompt.""" + + prompt: str + + +def _create_test_client_class( + generate_output: str = "test output", + class_name: str | None = None, +) -> type[Client]: + """Create a test client class with minimal implementation.""" + if class_name is None: + class_name = f"TestClient_{generate_output.replace(' ', '_')}" + + class TestClientClass(Client): + """Test client implementation.""" + + @classmethod + def parameter_mappers(cls) -> list[ParameterMapper]: + return [] + + def _init_request(self, inputs: Input) -> dict[str, Any]: + prompt = getattr(inputs, "prompt", "test prompt") + return {"prompt": prompt} + + def _parse_usage(self, response_data: dict[str, Any]) -> Usage: + return Usage() + + def _parse_content( + self, response_data: dict[str, Any], **parameters: Unpack[Parameters] + ) -> Any: # noqa: ANN401 + return response_data.get("content", "test content") + + async def generate(self, **parameters: Unpack[Parameters]) -> Output: + return Output(content=generate_output) + + TestClientClass.__name__ = class_name + return TestClientClass + + +def _create_test_mapper( + param_name: StrEnum, + map_key: str | None = None, +) -> ParameterMapper: + """Create a test parameter mapper instance.""" + actual_map_key = map_key if map_key is not None else param_name.value + + class TestMapperClass(ParameterMapper): + """Test mapper implementation.""" + + name: StrEnum = param_name + + def map( + self, + request: dict[str, Any], + value: Any, # noqa: ANN401 + model: Model, + ) -> dict[str, Any]: + if value is not None: + request[actual_map_key] = value + return request + + def parse_output(self, content: object, value: object | None) -> object: + return content + + return TestMapperClass() + + +def _create_transform_mapper( + param_name: StrEnum, + map_key: str | None = None, +) -> ParameterMapper: + """Create a test parameter mapper that transforms output.""" + actual_map_key = map_key if map_key is not None else param_name.value + + class TransformMapperClass(ParameterMapper): + """Test mapper with output transformation.""" + + name: StrEnum = param_name + + def map( + self, + request: dict[str, Any], + value: Any, # noqa: ANN401 + model: Model, + ) -> dict[str, Any]: + if value is not None: + request[actual_map_key] = value + return request + + def parse_output(self, content: object, value: object | None) -> object: + if value is not None: + return f"{content}_transformed_with_{value}" + return content + + return TransformMapperClass() + + +@pytest.fixture +def text_model() -> Model: + """Model that supports text generation.""" + return Model( + id="gpt-4", + provider=Provider.OPENAI, + capabilities={Capability.TEXT_GENERATION}, + display_name="GPT-4", + ) + + +@pytest.fixture +def multimodal_model() -> Model: + """Model that supports both text and image capabilities.""" + return Model( + id="gpt-4-vision", + provider=Provider.OPENAI, + capabilities={Capability.TEXT_GENERATION, Capability.IMAGE_GENERATION}, + display_name="GPT-4 Vision", + ) + + +@pytest.fixture +def api_key() -> str: + """Test API key.""" + return "sk-test123456789" + + +class ConcreteClient(Client): + """Concrete implementation for testing Client behavior.""" + + @classmethod + def parameter_mappers(cls) -> list[ParameterMapper]: + return [] + + def _init_request(self, inputs: Input) -> dict[str, Any]: + prompt = getattr(inputs, "prompt", "test prompt") + return {"prompt": prompt, "model": self.model.id} + + def _parse_usage(self, response_data: dict[str, Any]) -> Usage: + return Usage() + + def _parse_content( + self, response_data: dict[str, Any], **parameters: Unpack[Parameters] + ) -> Any: # noqa: ANN401 + return response_data.get("content", "test content") + + async def generate(self, **parameters: Unpack[Parameters]) -> Output: + return Output(content="test output") + + +class TestClientValidation: + """Test Client critical validation behaviors.""" + + @pytest.mark.smoke + def test_successful_creation_with_compatible_capability( + self, text_model: Model, api_key: str + ) -> None: + """Client accepts model that supports the required capability.""" + # Arrange & Act + client = ConcreteClient( + model=text_model, + provider=text_model.provider, + capability=Capability.TEXT_GENERATION, + api_key=SecretStr(api_key), + ) + + # Assert + assert client.model == text_model + assert client.capability == Capability.TEXT_GENERATION + + def test_validation_failure_with_incompatible_capability( + self, text_model: Model, api_key: str + ) -> None: + """Client rejects model that lacks required capability.""" + # Arrange & Act & Assert + with pytest.raises( + ValidationError, + match=r"Model 'gpt-4' does not support capability image_generation", + ): + ConcreteClient( + model=text_model, + provider=text_model.provider, + capability=Capability.IMAGE_GENERATION, # Model doesn't support this + api_key=SecretStr(api_key), + ) + + @pytest.mark.parametrize( + "capability,description", + [ + (Capability.TEXT_GENERATION, "text capability from multimodal model"), + (Capability.IMAGE_GENERATION, "image capability from multimodal model"), + ], + ids=["text_capability", "image_capability"], + ) + def test_validation_success_with_supported_capabilities( + self, + multimodal_model: Model, + api_key: str, + capability: Capability, + description: str, + ) -> None: + """Client accepts model that supports requested capability.""" + # Arrange & Act + client = ConcreteClient( + model=multimodal_model, + provider=multimodal_model.provider, + capability=capability, + api_key=SecretStr(api_key), + ) + + # Assert + assert client.model == multimodal_model + assert client.capability == capability + + def test_validation_fails_with_model_lacking_any_capabilities( + self, api_key: str + ) -> None: + """Client rejects models with empty capability set.""" + # Arrange + empty_model = Model( + id="broken-model", + provider=Provider.OPENAI, + capabilities=set(), # No capabilities + display_name="Broken Model", + ) + + # Act & Assert + with pytest.raises( + ValidationError, + match=r"Model 'broken-model' does not support capability text_generation", + ): + ConcreteClient( + model=empty_model, + provider=empty_model.provider, + capability=Capability.TEXT_GENERATION, + api_key=SecretStr(api_key), + ) + + +class TestClientRegistry: + """Test client registry functions - register_client and get_client_class.""" + + @pytest.fixture(autouse=True) + def clear_registry(self) -> Generator[None, None, None]: + """Clear the client registry before each test to ensure isolation.""" + # Arrange - Store original state and clear registry + original_clients = _clients.copy() + _clients.clear() + + yield + + # Cleanup - Restore original state + _clients.clear() + _clients.update(original_clients) + + @pytest.mark.smoke + def test_register_and_retrieve_client_success(self) -> None: + """Registry stores and retrieves client classes correctly.""" + # Arrange + capability = Capability.TEXT_GENERATION + provider = Provider.OPENAI + + # Act + register_client(capability, provider, ConcreteClient) + retrieved_class = get_client_class(capability, provider) + + # Assert + assert retrieved_class is ConcreteClient + + def test_get_client_class_raises_for_unregistered_capability(self) -> None: + """get_client_class raises NotImplementedError for unregistered capabilities.""" + # Arrange + unregistered_capability = Capability.IMAGE_GENERATION + provider = Provider.OPENAI + + # Act & Assert + with pytest.raises( + NotImplementedError, match=r"No client registered for image_generation" + ): + get_client_class(unregistered_capability, provider) + + def test_register_client_overwrites_previous_registration(self) -> None: + """Registering a new client for existing capability overwrites the previous one.""" + # Arrange + capability = Capability.TEXT_GENERATION + provider = Provider.OPENAI + + FirstClient = _create_test_client_class("first client", "FirstClient") + SecondClient = _create_test_client_class("second client", "SecondClient") + + # Act + register_client(capability, provider, FirstClient) + register_client(capability, provider, SecondClient) # Overwrite + retrieved_class = get_client_class(capability, provider) + + # Assert + assert retrieved_class is SecondClient + + def test_registry_isolation_between_different_capabilities(self) -> None: + """Different capabilities stored independently in the registry.""" + # Arrange + text_capability = Capability.TEXT_GENERATION + image_capability = Capability.IMAGE_GENERATION + provider = Provider.OPENAI + + TextClient = _create_test_client_class("text output", "TextClient") + ImageClient = _create_test_client_class("image output", "ImageClient") + + # Act + register_client(text_capability, provider, TextClient) + register_client(image_capability, provider, ImageClient) + + # Assert + assert get_client_class(text_capability, provider) is TextClient + assert get_client_class(image_capability, provider) is ImageClient + + @pytest.mark.parametrize( + "missing_capability,provider,expected_capability_str,expected_provider_str", + [ + ( + Capability.IMAGE_GENERATION, + Provider.ANTHROPIC, + "image_generation", + "anthropic", + ), + ( + Capability.VIDEO_GENERATION, + Provider.OPENAI, + "video_generation", + "openai", + ), + ], + ids=["image_anthropic", "video_openai"], + ) + def test_exception_message_includes_capability_and_provider( + self, + missing_capability: Capability, + provider: Provider, + expected_capability_str: str, + expected_provider_str: str, + ) -> None: + """NotImplementedError includes both capability and provider for debugging.""" + # Arrange & Act & Assert + with pytest.raises(NotImplementedError) as exc_info: + get_client_class(missing_capability, provider) + + # Assert both parts in error message + error_msg = str(exc_info.value) + assert expected_capability_str in error_msg + assert expected_provider_str in error_msg + + +class TestClientRequestBuilding: + """Test Client._build_request parameter mapping logic.""" + + @pytest.mark.smoke + def test_build_request_applies_parameter_mappers_correctly( + self, text_model: Model, api_key: str + ) -> None: + """_build_request applies all parameter mappers in sequence.""" + + # Arrange + class ClientWithMapper(ConcreteClient): + """Client with custom parameter mapper.""" + + @classmethod + def parameter_mappers(cls) -> list[ParameterMapper]: + """Return test mapper.""" + return [_create_test_mapper(ParamEnum.TEST_PARAM)] + + client = ClientWithMapper( + model=text_model, + provider=text_model.provider, + capability=Capability.TEXT_GENERATION, + api_key=SecretStr(api_key), + ) + + inputs = _TestInput(prompt="test prompt") + + # Act + request = client._build_request(inputs, test_param="mapped_value") # type: ignore[call-arg] + + # Assert + assert request["prompt"] == "test prompt" + assert request["test_param"] == "mapped_value" + + def test_build_request_with_multiple_mappers( + self, text_model: Model, api_key: str + ) -> None: + """_build_request applies multiple parameter mappers in order.""" + + # Arrange + class ClientWithMultipleMappers(ConcreteClient): + """Client with multiple parameter mappers.""" + + @classmethod + def parameter_mappers(cls) -> list[ParameterMapper]: + """Return multiple test mappers.""" + return [ + _create_test_mapper(ParamEnum.FIRST_PARAM), + _create_test_mapper(ParamEnum.SECOND_PARAM), + ] + + client = ClientWithMultipleMappers( + model=text_model, + provider=text_model.provider, + capability=Capability.TEXT_GENERATION, + api_key=SecretStr(api_key), + ) + + inputs = _TestInput(prompt="test prompt") + + # Act + request = client._build_request( + inputs, first_param="first", second_param="second" + ) # type: ignore[call-arg] + + # Assert + assert request["first_param"] == "first" + assert request["second_param"] == "second" + + @pytest.mark.parametrize( + "param_value,expected_output", + [ + ("test_value", "original content_transformed_with_test_value"), + (None, "original content"), + ], + ids=["with_value", "with_none"], + ) + def test_transform_output_applies_mappers( + self, + text_model: Model, + api_key: str, + param_value: str | None, + expected_output: str, + ) -> None: + """_transform_output applies parameter mapper output transformations.""" + + # Arrange + class ClientWithTransformMapper(ConcreteClient): + """Client with output transformation mapper.""" + + @classmethod + def parameter_mappers(cls) -> list[ParameterMapper]: + """Return transform mapper.""" + return [_create_transform_mapper(ParamEnum.TEST_PARAM)] + + client = ClientWithTransformMapper( + model=text_model, + provider=text_model.provider, + capability=Capability.TEXT_GENERATION, + api_key=SecretStr(api_key), + ) + + original_content = "original content" + + # Act + kwargs = {"test_param": param_value} if param_value is not None else {} + transformed = client._transform_output(original_content, **kwargs) + + # Assert + assert transformed == expected_output + + +class TestClientStreaming: + """Test Client.stream default behavior.""" + + def test_stream_raises_not_implemented_with_descriptive_error( + self, text_model: Model, api_key: str + ) -> None: + """stream() raises NotImplementedError with capability and provider info.""" + # Arrange + client = ConcreteClient( + model=text_model, + provider=text_model.provider, + capability=Capability.TEXT_GENERATION, + api_key=SecretStr(api_key), + ) + + # Act & Assert + with pytest.raises(NotImplementedError) as exc_info: + client.stream() + + # Verify error message contains all debugging info + error_msg = str(exc_info.value) + assert "Streaming not supported" in error_msg + assert "text_generation" in error_msg + assert "openai" in error_msg diff --git a/tests/unit_tests/test_constraints.py b/tests/unit_tests/test_constraints.py new file mode 100644 index 00000000..90f039a3 --- /dev/null +++ b/tests/unit_tests/test_constraints.py @@ -0,0 +1,337 @@ +"""Tests for constraint validation models.""" + +import pytest + +from celeste.constraints import Bool, Choice, Float, Int, Pattern, Range, Str + + +class TestChoice: + """Test Choice constraint validation.""" + + @pytest.mark.smoke + def test_validates_value_in_options(self) -> None: + """Test that valid choice passes validation.""" + constraint = Choice[str](options=["a", "b", "c"]) + + result = constraint("b") + + assert result == "b" + + def test_rejects_value_not_in_options(self) -> None: + """Test that invalid choice raises ValueError.""" + constraint = Choice[str](options=["a", "b", "c"]) + + with pytest.raises( + ValueError, match=r"Must be one of \['a', 'b', 'c'\], got 'd'" + ): + constraint("d") + + def test_works_with_numeric_types(self) -> None: + """Test Choice works with int/float options.""" + constraint = Choice[int](options=[1, 2, 3]) + + result = constraint(2) + + assert result == 2 + + def test_rejects_empty_options_list(self) -> None: + """Test Choice construction fails with empty options.""" + with pytest.raises(ValueError): + Choice[str](options=[]) + + +class TestRange: + """Test Range constraint validation.""" + + @pytest.mark.smoke + def test_validates_value_within_range(self) -> None: + """Test that value within bounds passes validation.""" + constraint = Range(min=0.0, max=1.0) + + result = constraint(0.5) + + assert result == 0.5 + + def test_validates_boundary_values(self) -> None: + """Test that min/max boundary values are inclusive.""" + constraint = Range(min=0, max=10) + + assert constraint(0) == 0 + assert constraint(10) == 10 + + def test_rejects_value_below_min(self) -> None: + """Test that value below min raises ValueError.""" + constraint = Range(min=0, max=10) + + with pytest.raises(ValueError, match=r"Must be between 0 and 10, got -1"): + constraint(-1) + + def test_rejects_value_above_max(self) -> None: + """Test that value above max raises ValueError.""" + constraint = Range(min=0, max=10) + + with pytest.raises(ValueError, match=r"Must be between 0 and 10, got 11"): + constraint(11) + + def test_rejects_non_numeric_value(self) -> None: + """Test that non-numeric value raises TypeError.""" + constraint = Range(min=0, max=10) + + with pytest.raises(TypeError, match=r"Must be numeric, got str"): + constraint("5") # type: ignore[arg-type] + + def test_accepts_both_int_and_float(self) -> None: + """Test Range accepts both int and float values.""" + constraint = Range(min=0.0, max=10.0) + + assert constraint(5) == 5 # int + assert constraint(5.5) == 5.5 # float + + def test_validates_value_with_step(self) -> None: + """Test that value at valid step increment passes.""" + constraint = Range(min=0, max=10, step=2) + + assert constraint(0) == 0 # min + assert constraint(2) == 2 + assert constraint(4) == 4 + assert constraint(10) == 10 # max + + def test_rejects_value_not_on_step(self) -> None: + """Test that value not on step increment raises ValueError.""" + constraint = Range(min=0, max=10, step=2) + + with pytest.raises( + ValueError, + match=r"Value must match step 2(\.0)?. Nearest valid: 2(\.0)? or 4(\.0)?, got 3", + ): + constraint(3) + + def test_validates_float_step(self) -> None: + """Test step validation with float increments.""" + constraint = Range(min=0.0, max=1.0, step=0.25) + + assert constraint(0.0) == 0.0 + assert constraint(0.25) == 0.25 + assert constraint(0.5) == 0.5 + assert constraint(0.75) == 0.75 + assert constraint(1.0) == 1.0 + + def test_step_validation_with_non_zero_min(self) -> None: + """Test step validation calculates offset from min correctly.""" + constraint = Range(min=5, max=15, step=3) + + assert constraint(5) == 5 # min + assert constraint(8) == 8 # min + 3 + assert constraint(11) == 11 # min + 6 + assert constraint(14) == 14 # min + 9 + + with pytest.raises( + ValueError, + match=r"Value must match step 3(\.0)?. Nearest valid: 5(\.0)? or 8(\.0)?, got 7", + ): + constraint(7) + + def test_step_validation_handles_float_precision(self) -> None: + """Test step validation handles floating-point precision issues.""" + constraint = Range(min=0.0, max=2.0, step=0.1) + + # These should all pass despite potential float precision issues + assert constraint(0.0) == 0.0 + assert constraint(0.1) == 0.1 + assert constraint(0.7) == 0.7 + assert constraint(1.0) == 1.0 + assert constraint(2.0) == 2.0 + + def test_validates_value_near_step_within_epsilon(self) -> None: + """Range validates values within epsilon tolerance of valid step.""" + constraint = Range(min=0.0, max=10.0, step=0.1) + + # Test values that might have floating-point precision issues + # 0.1 + 0.1 + 0.1 might be 0.30000000000000004 due to float representation + result = constraint(0.1 + 0.1 + 0.1) # Should be ~0.3 + assert result == pytest.approx(0.3) + + # Test another precision edge case + result2 = constraint(0.7) # Should pass within epsilon + assert result2 == pytest.approx(0.7) + + +class TestPattern: + """Test Pattern constraint validation.""" + + @pytest.mark.smoke + def test_validates_matching_pattern(self) -> None: + """Test that string matching pattern passes validation.""" + constraint = Pattern(pattern=r"^\d{3}-\d{4}$") + + result = constraint("123-4567") + + assert result == "123-4567" + + def test_rejects_non_matching_pattern(self) -> None: + """Test that non-matching pattern raises ValueError.""" + constraint = Pattern(pattern=r"^\d{3}-\d{4}$") + + with pytest.raises(ValueError, match=r"Must match pattern"): + constraint("abc-defg") + + def test_rejects_non_string_value(self) -> None: + """Test that non-string value raises TypeError.""" + constraint = Pattern(pattern=r"^\d+$") + + with pytest.raises(TypeError, match=r"Must be string, got int"): + constraint(123) # type: ignore[arg-type] + + def test_validates_complex_regex_patterns(self) -> None: + """Test Pattern works with complex regex.""" + # Email-like pattern + constraint = Pattern(pattern=r"^[a-z]+@[a-z]+\.[a-z]+$") + + result = constraint("user@domain.com") + + assert result == "user@domain.com" + + +class TestStr: + """Test Str constraint validation.""" + + @pytest.mark.smoke + def test_validates_string_without_length_constraints(self) -> None: + """Test that any string passes when no length constraints set.""" + constraint = Str() + + result = constraint("any string") + + assert result == "any string" + + def test_validates_string_within_length_bounds(self) -> None: + """Test string within min/max length passes.""" + constraint = Str(min_length=2, max_length=10) + + result = constraint("valid") + + assert result == "valid" + + def test_rejects_string_below_min_length(self) -> None: + """Test string shorter than min_length raises ValueError.""" + constraint = Str(min_length=5) + + with pytest.raises(ValueError, match=r"String too short \(min 5\), got 3"): + constraint("abc") + + def test_rejects_string_above_max_length(self) -> None: + """Test string longer than max_length raises ValueError.""" + constraint = Str(max_length=5) + + with pytest.raises(ValueError, match=r"String too long \(max 5\), got 10"): + constraint("too long!!") + + def test_rejects_non_string_value(self) -> None: + """Test non-string value raises TypeError.""" + constraint = Str() + + with pytest.raises(TypeError, match=r"Must be string, got int"): + constraint(123) # type: ignore[arg-type] + + def test_validates_boundary_lengths(self) -> None: + """Test exact min/max length strings are valid.""" + constraint = Str(min_length=3, max_length=5) + + assert constraint("abc") == "abc" # min + assert constraint("abcde") == "abcde" # max + + +class TestInt: + """Test Int constraint validation.""" + + @pytest.mark.smoke + def test_validates_integer_value(self) -> None: + """Test that integer passes validation.""" + constraint = Int() + + result = constraint(42) + + assert result == 42 + + def test_rejects_float_value(self) -> None: + """Test that float raises TypeError.""" + constraint = Int() + + with pytest.raises(TypeError, match=r"Must be int, got float"): + constraint(42.0) # type: ignore[arg-type] + + def test_rejects_boolean_value(self) -> None: + """Test that bool raises TypeError despite isinstance(True, int).""" + constraint = Int() + + with pytest.raises(TypeError, match=r"Must be int, got bool"): + constraint(True) + + def test_rejects_string_value(self) -> None: + """Test that string raises TypeError.""" + constraint = Int() + + with pytest.raises(TypeError, match=r"Must be int, got str"): + constraint("42") # type: ignore[arg-type] + + +class TestFloat: + """Test Float constraint validation.""" + + @pytest.mark.smoke + def test_validates_float_value(self) -> None: + """Test that float passes validation.""" + constraint = Float() + + result = constraint(3.14) + + assert result == 3.14 + + def test_accepts_and_converts_int_to_float(self) -> None: + """Test that int is accepted and converted to float.""" + constraint = Float() + + result = constraint(42) + + assert result == 42.0 + assert isinstance(result, float) + + def test_rejects_boolean_value(self) -> None: + """Test that bool raises TypeError despite isinstance(True, int).""" + constraint = Float() + + with pytest.raises(TypeError, match=r"Must be float or int, got bool"): + constraint(True) + + def test_rejects_string_value(self) -> None: + """Test that string raises TypeError.""" + constraint = Float() + + with pytest.raises(TypeError, match=r"Must be float or int, got str"): + constraint("3.14") # type: ignore[arg-type] + + +class TestBool: + """Test Bool constraint validation.""" + + @pytest.mark.smoke + def test_validates_boolean_value(self) -> None: + """Test that bool passes validation.""" + constraint = Bool() + + assert constraint(True) is True + assert constraint(False) is False + + def test_rejects_int_value(self) -> None: + """Test that int raises TypeError (no implicit 0/1 conversion).""" + constraint = Bool() + + with pytest.raises(TypeError, match=r"Must be bool, got int"): + constraint(1) # type: ignore[arg-type] + + def test_rejects_string_value(self) -> None: + """Test that string raises TypeError.""" + constraint = Bool() + + with pytest.raises(TypeError, match=r"Must be bool, got str"): + constraint("true") # type: ignore[arg-type] diff --git a/tests/unit_tests/test_core.py b/tests/unit_tests/test_core.py new file mode 100644 index 00000000..f69132c6 --- /dev/null +++ b/tests/unit_tests/test_core.py @@ -0,0 +1,174 @@ +"""High-value tests for core enums - focusing on critical framework behavior.""" + +import json +from enum import Enum + +import pytest + +from celeste.core import Capability, Provider + + +class TestProvider: + """Test Provider enum critical behaviors.""" + + @pytest.mark.smoke + def test_provider_is_string_enum(self) -> None: + """Provider must be a string enum for API compatibility.""" + # Arrange & Act & Assert + assert issubclass(Provider, str) + assert issubclass(Provider, Enum) + # Critical: string comparison must work for API responses + assert Provider.OPENAI == "openai" # type: ignore[comparison-overlap] + + def test_provider_json_serialization(self) -> None: + """Provider must serialize to JSON without custom encoder.""" + # Arrange + providers = { + "primary": Provider.ANTHROPIC, + "fallback": Provider.OPENAI, + } + + # Act + json_str = json.dumps(providers) + loaded = json.loads(json_str) + + # Assert + assert loaded["primary"] == "anthropic" + assert loaded["fallback"] == "openai" + + @pytest.mark.parametrize( + "provider_str,expected", + [ + ("openai", Provider.OPENAI), + ("google", Provider.GOOGLE), + ("anthropic", Provider.ANTHROPIC), + ], + ) + def test_provider_from_string(self, provider_str: str, expected: Provider) -> None: + """Provider can be constructed from string (common API pattern).""" + # Act + provider = Provider(provider_str) + + # Assert + assert provider == expected + assert provider.value == provider_str + + def test_invalid_provider_raises(self) -> None: + """Invalid provider string raises ValueError with helpful message.""" + # Act & Assert + with pytest.raises(ValueError) as exc_info: + Provider("invalid_provider") + + # Verify error message mentions valid options (helps debugging) + assert "'invalid_provider' is not a valid Provider" in str(exc_info.value) + + def test_provider_case_sensitive(self) -> None: + """Provider comparison is case-sensitive (API consistency).""" + # Act & Assert + assert Provider.OPENAI == "openai" # type: ignore[comparison-overlap] + assert Provider.OPENAI != "OpenAI" # type: ignore[comparison-overlap] + assert Provider.OPENAI != "OPENAI" + + +class TestCapability: + """Test Capability string enum critical behaviors.""" + + @pytest.mark.smoke + def test_capability_is_str_enum(self) -> None: + """Capability must be a str Enum for string-based capabilities.""" + # Arrange & Act & Assert + assert issubclass(Capability, Enum) + assert issubclass(Capability, str) + # Can create capability from string + cap = Capability("text_generation") + assert cap == Capability.TEXT_GENERATION + + def test_capability_set_operations(self) -> None: + """Capabilities work with set operations (multi-modal support).""" + # Arrange + text_only = Capability.TEXT_GENERATION + image_only = Capability.IMAGE_GENERATION + multi_modal = {text_only, image_only} + + # Act & Assert - critical for capability checking + assert text_only in multi_modal + assert image_only in multi_modal + assert len(multi_modal) == 2 + + def test_capability_values_are_unique(self) -> None: + """Each capability has a unique string value.""" + # Arrange + text = Capability.TEXT_GENERATION + image = Capability.IMAGE_GENERATION + + # Act & Assert + assert text.value != image.value + assert text.value == "text_generation" + assert image.value == "image_generation" + + def test_capability_string_conversion(self) -> None: + """Capability converts from and to strings correctly.""" + # Arrange & Act + from_string = Capability("text_generation") + from_enum = Capability.TEXT_GENERATION + + # Assert + assert from_string == from_enum + assert from_enum.value == "text_generation" + assert from_enum == "text_generation" + + +class TestEnumImmutability: + """Test that enums cannot be modified at runtime.""" + + def test_cannot_modify_provider_value(self) -> None: + """Provider enum values are immutable (prevents bugs).""" + # Act & Assert + with pytest.raises(AttributeError): + Provider.OPENAI.value = "modified" # type: ignore[misc] + + def test_cannot_delete_provider(self) -> None: + """Cannot delete existing providers (prevents accidental removal).""" + # Act & Assert + with pytest.raises(AttributeError): + del Provider.OPENAI + + def test_cannot_modify_capability_value(self) -> None: + """Capability enum values are immutable.""" + # Act & Assert + with pytest.raises(AttributeError): + Capability.TEXT_GENERATION.value = "modified" # type: ignore[misc] + + +class TestEnumUsagePatterns: + """Test common usage patterns in the framework.""" + + def test_provider_in_collection(self) -> None: + """Provider works correctly in sets and dicts (common pattern).""" + # Arrange & Act + providers_set = {Provider.OPENAI, Provider.GOOGLE, Provider.OPENAI} + provider_config = { + Provider.OPENAI: {"model": "gpt-4"}, + Provider.GOOGLE: {"model": "gemini"}, + } + + # Assert + assert len(providers_set) == 2 # Deduplication works + assert provider_config[Provider.OPENAI]["model"] == "gpt-4" + + def test_capability_empty_set(self) -> None: + """Can create empty capability set (no capabilities).""" + # Arrange + no_caps: set[Capability] = set() # Empty set of capabilities + + # Act & Assert + assert Capability.TEXT_GENERATION not in no_caps + assert Capability.IMAGE_GENERATION not in no_caps + assert not bool(no_caps) # Empty set is falsy + + @pytest.mark.parametrize("provider", list(Provider)) + def test_all_providers_are_lowercase(self, provider: Provider) -> None: + """All provider values are lowercase (API convention).""" + # Assert + assert provider.value == provider.value.lower() + assert provider.value.replace("_", "").isalpha() # Only letters diff --git a/tests/unit_tests/test_credentials.py b/tests/unit_tests/test_credentials.py new file mode 100644 index 00000000..573f0fa3 --- /dev/null +++ b/tests/unit_tests/test_credentials.py @@ -0,0 +1,370 @@ +"""Unit tests for credentials.py following 2025 best practices.""" + +import os +from collections.abc import Generator +from unittest.mock import patch + +import pytest +from pydantic import SecretStr + +from celeste.core import Provider +from celeste.credentials import PROVIDER_CREDENTIAL_MAP, Credentials + +# Single source of truth for environment variable names +ENV_VAR_NAMES = [ + "OPENAI_API_KEY", + "ANTHROPIC_API_KEY", + "GOOGLE_API_KEY", + "MISTRAL_API_KEY", + "HUGGINGFACE_TOKEN", + "STABILITYAI_API_KEY", + "REPLICATE_API_TOKEN", + "COHERE_API_KEY", + "XAI_API_KEY", + "LUMA_API_KEY", + "TOPAZLABS_API_KEY", +] + + +@pytest.fixture(autouse=True) +def clean_environment(monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]: + """Clear all environment variables before each test to ensure isolation. + + This prevents tests from failing when .env files load API keys into the + environment. Uses a two-part approach: + + 1. monkeypatch.setattr: Disables pydantic-settings .env file loading by + setting env_file=None. This prevents Credentials from reading .env + during initialization. + + 2. patch.dict(os.environ, clear=True): Clears all existing environment + variables so tests start with a clean slate. Required because .env + files may have already loaded vars into os.environ before test execution. + + Both are necessary: monkeypatch prevents future .env reads, patch.dict + clears existing state. + """ + # Disable .env file loading for Credentials class during tests + monkeypatch.setattr( + Credentials, + "model_config", + { + "env_file": None, + "env_file_encoding": "utf-8", + "case_sensitive": False, + "extra": "ignore", + }, + ) + + # Clear all environment variables + with patch.dict(os.environ, clear=True): + yield + + +class TestCredentialsLoading: + """Test loading credentials from environment.""" + + @pytest.mark.smoke + def test_load_from_env_single_provider( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Test loading a single provider credential from environment.""" + # Arrange + monkeypatch.setenv("OPENAI_API_KEY", "test-key") + + # Act + creds = Credentials() # type: ignore[call-arg] + + # Assert + assert creds.openai_api_key is not None + assert creds.openai_api_key.get_secret_value() == "test-key" + + def test_load_multiple_providers(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Test loading multiple provider credentials.""" + # Arrange + monkeypatch.setenv("OPENAI_API_KEY", "openai-key") + monkeypatch.setenv("ANTHROPIC_API_KEY", "anthropic-key") + + # Act + creds = Credentials() # type: ignore[call-arg] + + # Assert + assert creds.openai_api_key is not None + assert creds.anthropic_api_key is not None + assert creds.openai_api_key.get_secret_value() == "openai-key" + assert creds.anthropic_api_key.get_secret_value() == "anthropic-key" + assert creds.google_api_key is None + + def test_empty_credentials(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Test when no credentials are set.""" + # Arrange - clear any existing env vars + for env_var in ENV_VAR_NAMES: + monkeypatch.delenv(env_var, raising=False) + + # Act + creds = Credentials() # type: ignore[call-arg] + + # Assert - verify ALL credential fields are None + for _provider, field_name in PROVIDER_CREDENTIAL_MAP.items(): + assert getattr(creds, field_name) is None, ( + f"{field_name} should be None when no env vars set" + ) + + +class TestGetCredentials: + """Test get_credentials method following AAA pattern.""" + + @pytest.mark.smoke + def test_get_existing_credential(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Test retrieving an existing credential.""" + # Arrange + monkeypatch.setenv("OPENAI_API_KEY", "test-key") + creds = Credentials() # type: ignore[call-arg] + + # Act + result = creds.get_credentials(Provider.OPENAI) + + # Assert + assert isinstance(result, SecretStr) + assert result.get_secret_value() == "test-key" + + def test_get_missing_credential_raises( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Test that missing credentials raise ValueError.""" + # Arrange + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + creds = Credentials() # type: ignore[call-arg] + + # Act & Assert + with pytest.raises(ValueError, match="no credentials configured"): + creds.get_credentials(Provider.OPENAI) + + @pytest.mark.parametrize( + "provider,env_var,value", + [ + (Provider.OPENAI, "OPENAI_API_KEY", "openai-test"), + (Provider.ANTHROPIC, "ANTHROPIC_API_KEY", "anthropic-test"), + (Provider.GOOGLE, "GOOGLE_API_KEY", "google-test"), + (Provider.MISTRAL, "MISTRAL_API_KEY", "mistral-test"), + ], + ) + def test_get_credentials_parametrized( + self, + monkeypatch: pytest.MonkeyPatch, + provider: Provider, + env_var: str, + value: str, + ) -> None: + """Test get_credentials for various providers.""" + # Arrange + monkeypatch.setenv(env_var, value) + creds = Credentials() # type: ignore[call-arg] + + # Act + result = creds.get_credentials(provider) + + # Assert + assert result.get_secret_value() == value + + +class TestHasCredential: + """Test has_credential method.""" + + @pytest.mark.parametrize( + "provider,env_var", + [ + (Provider.OPENAI, "OPENAI_API_KEY"), + (Provider.GOOGLE, "GOOGLE_API_KEY"), + (Provider.ANTHROPIC, "ANTHROPIC_API_KEY"), + ], + ) + def test_has_credential_when_set( + self, monkeypatch: pytest.MonkeyPatch, provider: Provider, env_var: str + ) -> None: + """Test has_credential returns True when credential exists.""" + # Arrange + monkeypatch.setenv(env_var, "test-value") + creds = Credentials() # type: ignore[call-arg] + + # Act + result = creds.has_credential(provider) + + # Assert + assert result is True + + def test_has_credential_when_missing(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Test has_credential returns False when missing.""" + # Arrange + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + creds = Credentials() # type: ignore[call-arg] + + # Act & Assert + assert creds.has_credential(Provider.OPENAI) is False + + +class TestListAvailableProviders: + """Test list_available_providers method.""" + + @pytest.mark.smoke + def test_list_with_mixed_providers(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Test listing returns only providers with credentials.""" + # Arrange + monkeypatch.setenv("OPENAI_API_KEY", "test1") + monkeypatch.setenv("GOOGLE_API_KEY", "test2") + monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False) + creds = Credentials() # type: ignore[call-arg] + + # Act + providers = creds.list_available_providers() + + # Assert + assert Provider.OPENAI in providers + assert Provider.GOOGLE in providers + assert Provider.ANTHROPIC not in providers + assert len(providers) == 2 + + def test_empty_list_when_no_credentials( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Test empty list when no credentials configured.""" + # Arrange - clear all credential env vars + for env_var in ENV_VAR_NAMES: + monkeypatch.delenv(env_var, raising=False) + creds = Credentials() # type: ignore[call-arg] + + # Act + result = creds.list_available_providers() + + # Assert + assert result == [] + + +class TestCredentialSecurity: + """Test security features of credentials.""" + + @pytest.mark.smoke + def test_secret_str_hides_value(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Test SecretStr doesn't expose values in string representation.""" + # Arrange + secret_value = "super-secret-api-key" # nosec B105 + monkeypatch.setenv("OPENAI_API_KEY", secret_value) + creds = Credentials() # type: ignore[call-arg] + + # Act + credential = creds.get_credentials(Provider.OPENAI) + + # Assert + assert secret_value not in str(credential) + assert secret_value not in repr(credential) + assert credential.get_secret_value() == secret_value + + def test_model_dump_masks_secrets(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Test model_dump masks secret values.""" + # Arrange + monkeypatch.setenv("OPENAI_API_KEY", "secret-key") + creds = Credentials() # type: ignore[call-arg] + + # Act + dumped = creds.model_dump() + + # Assert + assert isinstance(dumped["openai_api_key"], SecretStr) + assert str(dumped["openai_api_key"]) == "**********" + assert "secret-key" not in str(dumped) + + def test_credentials_safe_for_logging( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ) -> None: + """Test that credentials object is safe to log without exposing secrets. + + This test prevents a critical security issue where developers might + accidentally log the entire credentials object, potentially exposing + API keys in production logs. + """ + import logging + + # Arrange + secret_value = "secret-key-12345" # nosec B105 + monkeypatch.setenv("OPENAI_API_KEY", secret_value) + monkeypatch.setenv("ANTHROPIC_API_KEY", "another-secret-xyz") + creds = Credentials() # type: ignore[call-arg] + + # Act - simulate common logging scenarios + logger = logging.getLogger(__name__) + caplog.set_level(logging.DEBUG) + + logger.info(f"Credentials initialized: {creds}") + logger.debug(f"Debug credentials repr: {creds!r}") + logger.error(f"Error with credentials: {creds!s}") + + # Assert - ensure no secrets appear in any log level + assert secret_value not in caplog.text + assert "another-secret-xyz" not in caplog.text + # Verify that some form of masking is present + assert "**********" in caplog.text or "SecretStr" in caplog.text + + +class TestProviderMapping: + """Test integrity of provider-credential mapping.""" + + def test_all_mapped_fields_exist(self) -> None: + """Test all fields in mapping exist on Credentials class.""" + # Arrange + creds = Credentials() # type: ignore[call-arg] + + # Act & Assert - verify each mapping points to a real field + for provider, field_name in PROVIDER_CREDENTIAL_MAP.items(): + assert hasattr(creds, field_name), ( + f"Missing field {field_name} for {provider}" + ) + + def test_all_providers_have_credential_mapping(self) -> None: + """Every Provider enum value has a corresponding entry in PROVIDER_CREDENTIAL_MAP.""" + # Get all providers that should have credentials + # Note: All providers in Provider enum require credentials + all_providers = list(Provider) + + # Verify each provider has a mapping + for provider in all_providers: + assert provider in PROVIDER_CREDENTIAL_MAP, ( + f"Provider {provider.value} missing from PROVIDER_CREDENTIAL_MAP" + ) + + +class TestEdgeCases: + """Test edge cases and error conditions.""" + + @pytest.mark.parametrize("value", ["", " ", "\t\n"]) + def test_whitespace_credentials( + self, monkeypatch: pytest.MonkeyPatch, value: str + ) -> None: + """Test handling of whitespace-only credentials. + + Note: Currently treats whitespace as 'present' credentials. + This is a deliberate choice - empty/whitespace values are still + considered as having credentials set (may be used for optional APIs). + """ + # Arrange + monkeypatch.setenv("OPENAI_API_KEY", value) + creds = Credentials() # type: ignore[call-arg] + + # Act & Assert + assert creds.has_credential(Provider.OPENAI) is True + credential = creds.get_credentials(Provider.OPENAI) + assert credential.get_secret_value() == value + + def test_special_characters_in_credentials( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Test handling of special characters in API keys.""" + # Arrange + special_key = "test!@#$%^&*()_+-=[]{}|;:',.<>?/~`" + monkeypatch.setenv("OPENAI_API_KEY", special_key) + creds = Credentials() # type: ignore[call-arg] + + # Act + credential = creds.get_credentials(Provider.OPENAI) + + # Assert + assert credential.get_secret_value() == special_key diff --git a/tests/unit_tests/test_http.py b/tests/unit_tests/test_http.py new file mode 100644 index 00000000..07ddaaf0 --- /dev/null +++ b/tests/unit_tests/test_http.py @@ -0,0 +1,833 @@ +"""High-value tests for HTTPClient - focusing on connection pooling and resource management.""" + +from collections.abc import AsyncIterator, Generator +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from celeste.core import Capability, Provider +from celeste.http import ( + HTTPClient, + clear_http_clients, + close_all_http_clients, + get_http_client, +) + + +@pytest.fixture +def mock_httpx_client() -> AsyncMock: + """Mock httpx.AsyncClient for testing HTTP operations.""" + client = AsyncMock(spec=httpx.AsyncClient) + client.post = AsyncMock(return_value=httpx.Response(200)) + client.get = AsyncMock(return_value=httpx.Response(200)) + client.aclose = AsyncMock() + return client + + +class TestHTTPClientLifecycle: + """Test HTTPClient initialization and cleanup behaviors.""" + + async def test_client_lazy_initialization(self) -> None: + """HTTPClient must not create httpx.AsyncClient until first request.""" + # Arrange + http_client = HTTPClient() + + # Assert - Client should not be initialized yet + assert http_client._client is None + + async def test_client_created_on_first_request( + self, mock_httpx_client: AsyncMock + ) -> None: + """HTTPClient must initialize httpx.AsyncClient on first request.""" + # Arrange + http_client = HTTPClient() + + # Act + with patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client): + await http_client.post( + url="https://api.example.com/test", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + + # Assert + assert http_client._client is not None + + async def test_client_reused_across_multiple_requests( + self, mock_httpx_client: AsyncMock + ) -> None: + """HTTPClient must reuse the same httpx.AsyncClient for multiple requests.""" + # Arrange + http_client = HTTPClient() + + # Act + with patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client): + await http_client.post( + url="https://api.example.com/test1", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value1"}, + ) + first_client = http_client._client + + await http_client.post( + url="https://api.example.com/test2", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value2"}, + ) + second_client = http_client._client + + # Assert - Same client instance must be reused + assert first_client is second_client + + async def test_aclose_sets_client_to_none( + self, mock_httpx_client: AsyncMock + ) -> None: + """HTTPClient.aclose() must properly cleanup and reset client state.""" + # Arrange + http_client = HTTPClient() + + # Act + with patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client): + await http_client.post( + url="https://api.example.com/test", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + assert http_client._client is not None + + await http_client.aclose() + + # Assert + assert http_client._client is None + mock_httpx_client.aclose.assert_called_once() + + async def test_aclose_handles_uninitialized_client(self) -> None: + """HTTPClient.aclose() must handle the case when client was never initialized.""" + # Arrange + http_client = HTTPClient() + + # Act & Assert - Should not raise any exceptions + await http_client.aclose() + assert http_client._client is None + + +class TestHTTPClientConnectionPooling: + """Test connection pool configuration behaviors.""" + + async def test_connection_limits_applied_to_httpx_client( + self, mock_httpx_client: AsyncMock + ) -> None: + """HTTPClient must configure httpx.AsyncClient with specified connection limits.""" + # Arrange + max_connections = 50 + max_keepalive = 25 + http_client = HTTPClient( + max_connections=max_connections, + max_keepalive_connections=max_keepalive, + ) + + # Act + with patch( + "celeste.http.httpx.AsyncClient", return_value=mock_httpx_client + ) as mock_constructor: + await http_client.post( + url="https://api.example.com/test", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + + # Assert - Verify AsyncClient was called with correct limits + mock_constructor.assert_called_once() + call_kwargs = mock_constructor.call_args[1] + limits = call_kwargs["limits"] + + assert limits.max_connections == max_connections + assert limits.max_keepalive_connections == max_keepalive + + async def test_default_connection_limits( + self, mock_httpx_client: AsyncMock + ) -> None: + """HTTPClient must use sensible default connection limits.""" + # Arrange + http_client = HTTPClient() # No explicit limits + + # Act + with patch( + "celeste.http.httpx.AsyncClient", return_value=mock_httpx_client + ) as mock_constructor: + await http_client.post( + url="https://api.example.com/test", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + + # Assert - Verify defaults are applied + mock_constructor.assert_called_once() + call_kwargs = mock_constructor.call_args[1] + limits = call_kwargs["limits"] + + assert limits.max_connections == 20 + assert limits.max_keepalive_connections == 10 + + +class TestHTTPClientRequestMethods: + """Test POST and GET request methods.""" + + async def test_post_request_with_all_parameters( + self, mock_httpx_client: AsyncMock + ) -> None: + """POST method must pass all parameters correctly to httpx.AsyncClient.""" + # Arrange + http_client = HTTPClient() + url = "https://api.example.com/generate" + headers = { + "Authorization": "Bearer sk-test", + "Content-Type": "application/json", + } + json_body = {"prompt": "Hello", "max_tokens": 100} + timeout = 30.0 + + mock_response = httpx.Response(200, json={"result": "success"}) + mock_httpx_client.post.return_value = mock_response + + # Act + with patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client): + response = await http_client.post( + url=url, + headers=headers, + json_body=json_body, + timeout=timeout, + ) + + # Assert + mock_httpx_client.post.assert_called_once_with( + url, + headers=headers, + json=json_body, + timeout=timeout, + ) + assert response.status_code == 200 + + async def test_get_request_with_all_parameters( + self, mock_httpx_client: AsyncMock + ) -> None: + """GET method must pass all parameters correctly to httpx.AsyncClient.""" + # Arrange + http_client = HTTPClient() + url = "https://api.example.com/models" + headers = {"Authorization": "Bearer sk-test"} + timeout = 15.0 + + mock_response = httpx.Response(200, json={"models": ["gpt-4"]}) + mock_httpx_client.get.return_value = mock_response + + # Act + with patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client): + response = await http_client.get( + url=url, + headers=headers, + timeout=timeout, + ) + + # Assert + mock_httpx_client.get.assert_called_once_with( + url, + headers=headers, + timeout=timeout, + follow_redirects=True, + ) + assert response.status_code == 200 + + async def test_post_propagates_httpx_errors( + self, mock_httpx_client: AsyncMock + ) -> None: + """POST method must propagate httpx.HTTPError to caller.""" + # Arrange + http_client = HTTPClient() + mock_httpx_client.post.side_effect = httpx.TimeoutException("Request timeout") + + # Act & Assert + with ( + patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client), + pytest.raises(httpx.TimeoutException, match="Request timeout"), + ): + await http_client.post( + url="https://api.example.com/test", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + + async def test_get_propagates_httpx_errors( + self, mock_httpx_client: AsyncMock + ) -> None: + """GET method must propagate httpx.HTTPError to caller.""" + # Arrange + http_client = HTTPClient() + mock_httpx_client.get.side_effect = httpx.ConnectError("Connection failed") + + # Act & Assert + with ( + patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client), + pytest.raises(httpx.ConnectError, match="Connection failed"), + ): + await http_client.get( + url="https://api.example.com/test", + headers={"Authorization": "Bearer test"}, + ) + + async def test_post_uses_default_timeout_when_not_specified( + self, mock_httpx_client: AsyncMock + ) -> None: + """POST method must use default timeout when none is specified.""" + # Arrange + http_client = HTTPClient() + + # Act + with patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client): + await http_client.post( + url="https://api.example.com/test", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + # Note: timeout parameter omitted + ) + + # Assert - Verify default timeout was used + mock_httpx_client.post.assert_called_once() + call_kwargs = mock_httpx_client.post.call_args[1] + assert call_kwargs["timeout"] == 60.0 + + async def test_get_uses_default_timeout_when_not_specified( + self, mock_httpx_client: AsyncMock + ) -> None: + """GET method must use default timeout when none is specified.""" + # Arrange + http_client = HTTPClient() + + # Act + with patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client): + await http_client.get( + url="https://api.example.com/test", + headers={"Authorization": "Bearer test"}, + # Note: timeout parameter omitted + ) + + # Assert - Verify default timeout was used + mock_httpx_client.get.assert_called_once() + call_kwargs = mock_httpx_client.get.call_args[1] + assert call_kwargs["timeout"] == 60.0 + + async def test_custom_timeout_passed_to_httpx( + self, mock_httpx_client: AsyncMock + ) -> None: + """HTTPClient passes custom timeout value to httpx.AsyncClient methods.""" + http_client = HTTPClient() + custom_timeout = 120.0 + + with patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client): + await http_client.post( + url="https://api.example.com/test", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + timeout=custom_timeout, + ) + + # Verify custom timeout was passed to httpx + mock_httpx_client.post.assert_called_once() + call_kwargs = mock_httpx_client.post.call_args[1] + assert call_kwargs["timeout"] == custom_timeout + + +class TestHTTPClientRegistry: + """Test get_http_client registry and singleton behavior.""" + + @pytest.fixture(autouse=True) + def clear_registry(self) -> Generator[None, None, None]: + """Clear HTTP client registry before each test to ensure isolation.""" + from celeste.http import _http_clients + + # Arrange - Store original state and clear registry + original_clients = _http_clients.copy() + _http_clients.clear() + + yield + + # Cleanup - Restore original state + _http_clients.clear() + _http_clients.update(original_clients) + + def test_get_http_client_returns_same_instance_for_same_key(self) -> None: + """get_http_client must return the same HTTPClient for identical provider/capability.""" + # Arrange + provider = Provider.OPENAI + capability = Capability.TEXT_GENERATION + + # Act + client1 = get_http_client(provider, capability) + client2 = get_http_client(provider, capability) + + # Assert - Must be the exact same instance (singleton behavior) + assert client1 is client2 + + def test_get_http_client_returns_different_instances_for_different_providers( + self, + ) -> None: + """get_http_client must return different HTTPClients for different providers.""" + # Arrange + capability = Capability.TEXT_GENERATION + + # Act + openai_client = get_http_client(Provider.OPENAI, capability) + anthropic_client = get_http_client(Provider.ANTHROPIC, capability) + + # Assert - Must be different instances + assert openai_client is not anthropic_client + + def test_get_http_client_returns_different_instances_for_different_capabilities( + self, + ) -> None: + """get_http_client must return different HTTPClients for different capabilities.""" + # Arrange + provider = Provider.OPENAI + + # Act + text_client = get_http_client(provider, Capability.TEXT_GENERATION) + image_client = get_http_client(provider, Capability.IMAGE_GENERATION) + + # Assert - Must be different instances + assert text_client is not image_client + + def test_registry_isolation_prevents_cross_contamination(self) -> None: + """Registry must maintain complete isolation between different provider/capability pairs.""" + # Arrange - Create clients for different combinations + openai_text = get_http_client(Provider.OPENAI, Capability.TEXT_GENERATION) + openai_image = get_http_client(Provider.OPENAI, Capability.IMAGE_GENERATION) + anthropic_text = get_http_client(Provider.ANTHROPIC, Capability.TEXT_GENERATION) + + # Act - Retrieve them again + openai_text_again = get_http_client(Provider.OPENAI, Capability.TEXT_GENERATION) + openai_image_again = get_http_client( + Provider.OPENAI, Capability.IMAGE_GENERATION + ) + anthropic_text_again = get_http_client( + Provider.ANTHROPIC, Capability.TEXT_GENERATION + ) + + # Assert - Same pairs return same instances, different pairs return different instances + assert openai_text is openai_text_again + assert openai_image is openai_image_again + assert anthropic_text is anthropic_text_again + + # Verify all three are distinct + assert openai_text is not openai_image + assert openai_text is not anthropic_text + assert openai_image is not anthropic_text + + +class TestHTTPClientCleanup: + """Test close_all_http_clients and clear_http_clients functionality.""" + + @pytest.fixture(autouse=True) + def clear_registry(self) -> Generator[None, None, None]: + """Clear HTTP client registry before each test to ensure isolation.""" + from celeste.http import _http_clients + + # Arrange - Store original state and clear registry + original_clients = _http_clients.copy() + _http_clients.clear() + + yield + + # Cleanup - Restore original state + _http_clients.clear() + _http_clients.update(original_clients) + + async def test_close_all_http_clients_closes_all_and_clears_registry( + self, mock_httpx_client: AsyncMock + ) -> None: + """close_all_http_clients must close all clients and clear the registry.""" + from celeste.http import _http_clients, get_http_client + + # Arrange - Create multiple clients + client1 = get_http_client(Provider.OPENAI, Capability.TEXT_GENERATION) + client2 = get_http_client(Provider.ANTHROPIC, Capability.TEXT_GENERATION) + + # Initialize both clients to create httpx.AsyncClient instances + with patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client): + await client1.post( + url="https://api.example.com/test1", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + await client2.post( + url="https://api.example.com/test2", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + + # Verify registry has clients and clients are initialized + assert len(_http_clients) == 2 + assert client1._client is not None + assert client2._client is not None + + # Act - Close all clients + await close_all_http_clients() + + # Assert - Registry should be empty AND clients should be reset + assert len(_http_clients) == 0 + assert client1._client is None + assert client2._client is None + + async def test_close_all_http_clients_calls_aclose_on_each_client( + self, mock_httpx_client: AsyncMock + ) -> None: + """close_all_http_clients must call aclose() on each httpx.AsyncClient.""" + from celeste.http import get_http_client + + # Arrange - Create and initialize multiple clients with separate mock instances + mock_client1 = AsyncMock(spec=httpx.AsyncClient) + mock_client1.post = AsyncMock(return_value=httpx.Response(200)) + mock_client1.aclose = AsyncMock() + + mock_client2 = AsyncMock(spec=httpx.AsyncClient) + mock_client2.post = AsyncMock(return_value=httpx.Response(200)) + mock_client2.aclose = AsyncMock() + + client1 = get_http_client(Provider.OPENAI, Capability.TEXT_GENERATION) + client2 = get_http_client(Provider.ANTHROPIC, Capability.TEXT_GENERATION) + + # Initialize clients with different mock instances + with patch("celeste.http.httpx.AsyncClient") as mock_constructor: + mock_constructor.side_effect = [mock_client1, mock_client2] + + await client1.post( + url="https://api.example.com/test1", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + await client2.post( + url="https://api.example.com/test2", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + + # Act - Close all clients + await close_all_http_clients() + + # Assert - Both mock clients should have aclose called + mock_client1.aclose.assert_called_once() + mock_client2.aclose.assert_called_once() + + async def test_clear_http_clients_clears_without_closing(self) -> None: + """clear_http_clients must clear registry without calling aclose().""" + from celeste.http import _http_clients, get_http_client + + # Arrange - Create multiple clients + mock_client1 = AsyncMock(spec=httpx.AsyncClient) + mock_client1.post = AsyncMock(return_value=httpx.Response(200)) + mock_client1.aclose = AsyncMock() + + client1 = get_http_client(Provider.OPENAI, Capability.TEXT_GENERATION) + + # Initialize client + with patch("celeste.http.httpx.AsyncClient", return_value=mock_client1): + await client1.post( + url="https://api.example.com/test", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + + # Verify registry has client + assert len(_http_clients) == 1 + + # Act - Clear without closing + clear_http_clients() + + # Assert - Registry cleared but aclose not called + assert len(_http_clients) == 0 + mock_client1.aclose.assert_not_called() + + async def test_close_all_handles_multiple_providers_and_capabilities( + self, + ) -> None: + """close_all_http_clients must handle multiple provider/capability combinations.""" + from celeste.http import _http_clients, get_http_client + + # Arrange - Create clients for different combinations + mock_client1 = AsyncMock(spec=httpx.AsyncClient) + mock_client1.post = AsyncMock(return_value=httpx.Response(200)) + mock_client1.aclose = AsyncMock() + + mock_client2 = AsyncMock(spec=httpx.AsyncClient) + mock_client2.post = AsyncMock(return_value=httpx.Response(200)) + mock_client2.aclose = AsyncMock() + + mock_client3 = AsyncMock(spec=httpx.AsyncClient) + mock_client3.post = AsyncMock(return_value=httpx.Response(200)) + mock_client3.aclose = AsyncMock() + + client1 = get_http_client(Provider.OPENAI, Capability.TEXT_GENERATION) + client2 = get_http_client(Provider.OPENAI, Capability.IMAGE_GENERATION) + client3 = get_http_client(Provider.ANTHROPIC, Capability.TEXT_GENERATION) + + # Initialize all clients + with patch("celeste.http.httpx.AsyncClient") as mock_constructor: + mock_constructor.side_effect = [mock_client1, mock_client2, mock_client3] + + await client1.post( + url="https://api.example.com/test1", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + await client2.post( + url="https://api.example.com/test2", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + await client3.post( + url="https://api.example.com/test3", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + + # Verify registry has all 3 clients + assert len(_http_clients) == 3 + + # Act - Close all + await close_all_http_clients() + + # Assert - All closed and registry empty + assert len(_http_clients) == 0 + mock_client1.aclose.assert_called_once() + mock_client2.aclose.assert_called_once() + mock_client3.aclose.assert_called_once() + + async def test_close_all_continues_despite_individual_failures(self) -> None: + """close_all_http_clients must continue closing all clients even if one fails.""" + from celeste.http import _http_clients, get_http_client + + # Arrange - Create clients where one will fail to close + mock_client1 = AsyncMock(spec=httpx.AsyncClient) + mock_client1.post = AsyncMock(return_value=httpx.Response(200)) + mock_client1.aclose = AsyncMock(side_effect=RuntimeError("Close failed")) + + mock_client2 = AsyncMock(spec=httpx.AsyncClient) + mock_client2.post = AsyncMock(return_value=httpx.Response(200)) + mock_client2.aclose = AsyncMock() + + client1 = get_http_client(Provider.OPENAI, Capability.TEXT_GENERATION) + client2 = get_http_client(Provider.ANTHROPIC, Capability.TEXT_GENERATION) + + # Initialize both clients + with patch("celeste.http.httpx.AsyncClient") as mock_constructor: + mock_constructor.side_effect = [mock_client1, mock_client2] + + await client1.post( + url="https://api.example.com/test1", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + await client2.post( + url="https://api.example.com/test2", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + + # Act - Should not raise, continues closing despite failure + await close_all_http_clients() + + # Assert - Both aclose calls were attempted and registry is cleared + mock_client1.aclose.assert_called_once() + mock_client2.aclose.assert_called_once() + assert len(_http_clients) == 0 + + +class TestHTTPClientContextManager: + """Test async context manager protocol for resource cleanup.""" + + async def test_context_manager_returns_self( + self, mock_httpx_client: AsyncMock + ) -> None: + """__aenter__ must return HTTPClient instance for use in async with block.""" + # Arrange + http_client = HTTPClient() + + # Act + async with http_client as client: + # Assert - Context manager returns the HTTPClient instance itself + assert client is http_client + + async def test_context_manager_calls_aclose_on_exit( + self, mock_httpx_client: AsyncMock + ) -> None: + """__aexit__ must call aclose() to cleanup connections on context exit.""" + # Arrange + http_client = HTTPClient() + + # Act - Initialize client and exit context + with patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client): + async with http_client: + await http_client.post( + url="https://api.example.com/test", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + # Verify client was initialized + assert http_client._client is not None + + # Assert - After exiting context, client should be closed + assert http_client._client is None + mock_httpx_client.aclose.assert_called_once() + + async def test_context_manager_cleanup_on_exception( + self, mock_httpx_client: AsyncMock + ) -> None: + """__aexit__ must cleanup connections even when exception occurs in context.""" + # Arrange + http_client = HTTPClient() + + # Act & Assert - Verify cleanup happens despite exception + with ( + patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client), + pytest.raises(ValueError, match="Test exception"), + ): + async with http_client: + await http_client.post( + url="https://api.example.com/test", + headers={"Authorization": "Bearer test"}, + json_body={"key": "value"}, + ) + raise ValueError("Test exception") + + # Assert - Client was closed despite exception + assert http_client._client is None + mock_httpx_client.aclose.assert_called_once() + + +class TestHTTPClientStreaming: + """Test Server-Sent Events streaming functionality.""" + + def _create_mock_sse(self, data: str) -> MagicMock: + """Create mock SSE event with data.""" + mock = MagicMock() + mock.data = data + return mock + + def _create_mock_event_source(self, events: list) -> MagicMock: + """Create mock SSE event source with events.""" + mock_source = MagicMock() + mock_source.aiter_sse = MagicMock(return_value=self._async_iter(events)) + mock_source.__aenter__ = AsyncMock(return_value=mock_source) + mock_source.__aexit__ = AsyncMock(return_value=False) + + # Mock response with headers for httpx_sse content-type check + mock_response = MagicMock() + mock_response.headers = MagicMock() + mock_response.headers.get = MagicMock(return_value="text/event-stream") + mock_source._response = mock_response + + return mock_source + + async def test_stream_post_yields_parsed_json_events( + self, mock_httpx_client: AsyncMock + ) -> None: + """stream_post must parse SSE events and yield all JSON.""" + # Arrange + http_client = HTTPClient() + events = [ + self._create_mock_sse('{"delta": "Hello"}'), + self._create_mock_sse('{"delta": " world"}'), + ] + mock_event_source = self._create_mock_event_source(events) + + # Act + with ( + patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client), + patch("celeste.http.aconnect_sse", return_value=mock_event_source), + ): + chunks = [ + chunk + async for chunk in http_client.stream_post( + url="https://api.example.com/stream", + headers={"Authorization": "Bearer test"}, + json_body={"prompt": "test"}, + ) + ] + + # Assert - All valid JSON events are parsed and yielded + assert chunks == [{"delta": "Hello"}, {"delta": " world"}] + + async def test_stream_post_raises_on_malformed_json( + self, mock_httpx_client: AsyncMock + ) -> None: + """stream_post must raise JSONDecodeError on malformed SSE events.""" + # Arrange + http_client = HTTPClient() + events = [ + self._create_mock_sse('{"valid": true}'), + self._create_mock_sse("{invalid json"), + self._create_mock_sse('{"valid": "also"}'), + ] + mock_event_source = self._create_mock_event_source(events) + + # Act + results = [] + with ( + patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client), + patch("celeste.http.aconnect_sse", return_value=mock_event_source), + ): + async for chunk in http_client.stream_post( + url="https://api.example.com/stream", + headers={"Authorization": "Bearer test"}, + json_body={"prompt": "test"}, + ): + results.append(chunk) + + # Assert - only valid JSON events are yielded + assert len(results) == 2 + assert results[0] == {"valid": True} + assert results[1] == {"valid": "also"} + + async def test_stream_post_passes_parameters_correctly( + self, mock_httpx_client: AsyncMock + ) -> None: + """stream_post must pass all parameters to aconnect_sse correctly.""" + # Arrange + http_client = HTTPClient() + url = "https://api.example.com/stream" + headers = {"Authorization": "Bearer sk-test", "X-Custom": "value"} + json_body = {"prompt": "test", "stream": True} + timeout = 120.0 + + mock_event_source = self._create_mock_event_source([]) # Empty stream + + # Act + with ( + patch("celeste.http.httpx.AsyncClient", return_value=mock_httpx_client), + patch( + "celeste.http.aconnect_sse", return_value=mock_event_source + ) as mock_sse, + ): + async for _ in http_client.stream_post( + url=url, + headers=headers, + json_body=json_body, + timeout=timeout, + ): + pass + + # Assert - Verify aconnect_sse called with correct parameters + mock_sse.assert_called_once_with( + mock_httpx_client, + "POST", + url, + json=json_body, + headers=headers, + timeout=timeout, + ) + + @staticmethod + async def _async_iter(items: list) -> AsyncIterator: + """Helper to create async iterator from list.""" + for item in items: + yield item diff --git a/tests/unit_tests/test_init.py b/tests/unit_tests/test_init.py new file mode 100644 index 00000000..c079079d --- /dev/null +++ b/tests/unit_tests/test_init.py @@ -0,0 +1,153 @@ +"""High-value tests for celeste.__init__ module.""" + +from unittest.mock import patch + +import pytest + +from celeste import Capability, Model, Provider, create_client + + +@pytest.fixture +def sample_models() -> list[Model]: + """Test models for various scenarios.""" + return [ + Model( + id="gpt-4", + provider=Provider.OPENAI, + capabilities={Capability.TEXT_GENERATION}, + display_name="GPT-4", + ), + Model( + id="claude-3", + provider=Provider.ANTHROPIC, + capabilities={Capability.TEXT_GENERATION}, + display_name="Claude 3", + ), + Model( + id="dall-e-3", + provider=Provider.OPENAI, + capabilities={Capability.IMAGE_GENERATION}, + display_name="DALL-E 3", + ), + ] + + +class TestCreateClient: + """Test the create_client factory function.""" + + def test_create_client_no_models_available_raises_error(self) -> None: + """Test that create_client raises ValueError when no models are available.""" + with patch("celeste.list_models", autospec=True) as mock_list_models: + # Arrange + mock_list_models.return_value = [] + + # Act & Assert + with pytest.raises( + ValueError, match=r"No model found for.*text_generation" + ): + create_client(capability=Capability.TEXT_GENERATION) + + def test_create_client_specific_model_not_found_raises_error(self) -> None: + """Test error when specific model/provider combination doesn't exist.""" + with patch("celeste.get_model", autospec=True) as mock_get_model: + # Arrange + mock_get_model.return_value = None + + # Act & Assert + with pytest.raises(ValueError, match=r"Model.*not found"): + create_client( + capability=Capability.TEXT_GENERATION, + provider=Provider.OPENAI, + model="nonexistent-model", + ) + + def test_create_client_uses_explicit_model_when_both_provided( + self, sample_models: list[Model] + ) -> None: + """Test that create_client uses get_model for explicit selection.""" + with patch("celeste.get_model", autospec=True) as mock_get_model: + # Arrange + mock_get_model.return_value = sample_models[1] # claude-3 + + # Act & Assert + with pytest.raises( + NotImplementedError + ): # _get_client_class not implemented + create_client( + capability=Capability.TEXT_GENERATION, + provider=Provider.ANTHROPIC, + model="claude-3", + ) + + mock_get_model.assert_called_once_with("claude-3", Provider.ANTHROPIC) + + def test_create_client_string_model_without_provider_raises_error(self) -> None: + """Test that string model ID without provider raises ValueError.""" + # Act & Assert + with pytest.raises( + ValueError, match="provider required when model is a string ID" + ): + create_client( + capability=Capability.TEXT_GENERATION, + model="some-model", # provider=None - should error + ) + + def test_create_client_filters_by_provider_when_specified( + self, sample_models: list[Model] + ) -> None: + """Test that provider filtering is applied when provider is specified.""" + with patch("celeste.list_models", autospec=True) as mock_list_models: + # Arrange + mock_list_models.return_value = [sample_models[1]] # claude-3 + + # Act - Should fail at _get_client_class but provider filtering should work + with pytest.raises( + NotImplementedError + ): # _get_client_class not implemented + create_client( + capability=Capability.TEXT_GENERATION, + provider=Provider.ANTHROPIC, + ) + + # Assert - verify provider was passed to list_models + mock_list_models.assert_called_once_with( + provider=Provider.ANTHROPIC, capability=Capability.TEXT_GENERATION + ) + + +class TestCreateClientIntegration: + """Test create_client integration with model selection.""" + + def test_model_selection_precedence(self, sample_models: list[Model]) -> None: + """Test that explicit model/provider takes precedence over auto-selection.""" + with ( + patch("celeste.list_models", autospec=True) as mock_list_models, + patch("celeste.get_model", autospec=True) as mock_get_model, + ): + # Arrange + explicit_model = sample_models[1] # claude-3 + auto_model = sample_models[0] # gpt-4 + + mock_get_model.return_value = explicit_model + mock_list_models.return_value = [auto_model] + + # Act - Should fail at _get_client_class but precedence should work + with pytest.raises(NotImplementedError): + create_client( + capability=Capability.TEXT_GENERATION, + provider=Provider.ANTHROPIC, + model="claude-3", + ) + + # Assert - Should use explicit path, not auto-selection + mock_get_model.assert_called_once_with("claude-3", Provider.ANTHROPIC) + mock_list_models.assert_not_called() # Should not try auto-selection + + def test_error_propagation_from_registry(self) -> None: + """Test that errors from registry functions propagate correctly.""" + # Test that registry errors bubble up properly + with patch("celeste.list_models", autospec=True) as mock_list_models: + mock_list_models.side_effect = ValueError("Registry error") + + with pytest.raises(ValueError, match="Registry error"): + create_client(capability=Capability.TEXT_GENERATION) diff --git a/tests/unit_tests/test_mime_types.py b/tests/unit_tests/test_mime_types.py new file mode 100644 index 00000000..f1114217 --- /dev/null +++ b/tests/unit_tests/test_mime_types.py @@ -0,0 +1,155 @@ +"""High-value tests for MIME type enums - focusing on inheritance and string behavior.""" + +import json +from enum import Enum + +import pytest + +from celeste.mime_types import AudioMimeType, ImageMimeType, MimeType, VideoMimeType + + +class TestMimeTypeInheritance: + """Test the MIME type enum inheritance structure.""" + + def test_all_mime_types_inherit_from_base(self) -> None: + """All specific MIME type enums should inherit from MimeType base class.""" + assert issubclass(ImageMimeType, MimeType) + assert issubclass(VideoMimeType, MimeType) + assert issubclass(AudioMimeType, MimeType) + + def test_mime_type_base_is_string_enum(self) -> None: + """MimeType should be both a string and Enum (StrEnum pattern).""" + assert issubclass(MimeType, str) + assert issubclass(MimeType, Enum) + + @pytest.mark.parametrize( + "mime_type,expected_value", + [ + (ImageMimeType.PNG, "image/png"), + (ImageMimeType.JPEG, "image/jpeg"), + (VideoMimeType.MP4, "video/mp4"), + (AudioMimeType.MP3, "audio/mpeg"), + (AudioMimeType.WAV, "audio/wav"), + ], + ) + def test_mime_types_equal_their_string_values( + self, mime_type: MimeType, expected_value: str + ) -> None: + """MIME types should equal their string values for API compatibility.""" + assert mime_type == expected_value + assert mime_type.value == expected_value + + +class TestImageMimeType: + """Test ImageMimeType specific values and behavior.""" + + def test_image_mime_type_values(self) -> None: + """ImageMimeType should have correct MIME type strings.""" + assert ImageMimeType.PNG.value == "image/png" + assert ImageMimeType.JPEG.value == "image/jpeg" + + def test_image_mime_type_members_exist(self) -> None: + """ImageMimeType should contain expected members.""" + members = list(ImageMimeType) + assert ImageMimeType.PNG in members + assert ImageMimeType.JPEG in members + + +class TestVideoMimeType: + """Test VideoMimeType specific values and behavior.""" + + def test_video_mime_type_values(self) -> None: + """VideoMimeType should have correct MIME type strings.""" + assert VideoMimeType.MP4.value == "video/mp4" + + def test_video_mime_type_members_exist(self) -> None: + """VideoMimeType should contain expected members.""" + members = list(VideoMimeType) + assert VideoMimeType.MP4 in members + + +class TestAudioMimeType: + """Test AudioMimeType specific values and behavior.""" + + def test_audio_mime_type_values(self) -> None: + """AudioMimeType should have correct MIME type strings.""" + assert AudioMimeType.MP3.value == "audio/mpeg" + assert AudioMimeType.WAV.value == "audio/wav" + + +class TestMimeTypeUsagePatterns: + """Test common usage patterns and edge cases.""" + + def test_mime_type_json_serialization(self) -> None: + """MIME types should serialize to JSON correctly.""" + # Critical for API responses + mime_dict = { + "image": ImageMimeType.JPEG, + "video": VideoMimeType.MP4, + "audio": AudioMimeType.MP3, + } + + # Should serialize without custom encoder + json_str = json.dumps(mime_dict) + loaded = json.loads(json_str) + + assert loaded["image"] == "image/jpeg" + assert loaded["video"] == "video/mp4" + assert loaded["audio"] == "audio/mpeg" + + def test_mime_type_membership_check(self) -> None: + """Test 'in' operator works with MIME type enums.""" + # Common pattern for validation + valid_image_types = [ImageMimeType.PNG, ImageMimeType.JPEG] + assert ImageMimeType.PNG in valid_image_types + assert VideoMimeType.MP4 not in valid_image_types # type: ignore[comparison-overlap] + + @pytest.mark.parametrize( + "mime_type,expected_category", + [ + (ImageMimeType.PNG, "image"), + (ImageMimeType.JPEG, "image"), + (VideoMimeType.MP4, "video"), + (AudioMimeType.MP3, "audio"), + (AudioMimeType.WAV, "audio"), + ], + ) + def test_mime_type_category_extraction( + self, mime_type: MimeType, expected_category: str + ) -> None: + """Test extracting category from MIME type string (common parsing need).""" + # Tests the string nature of our enums + category = mime_type.value.split("/")[0] + assert category == expected_category + + +class TestMimeTypeErrorCases: + """Test error handling and edge cases.""" + + def test_invalid_mime_type_comparison(self) -> None: + """Different MIME type categories should not be equal.""" + assert ImageMimeType.PNG != VideoMimeType.MP4 # type: ignore[comparison-overlap] + assert ImageMimeType.PNG != AudioMimeType.MP3 # type: ignore[comparison-overlap] + assert VideoMimeType.MP4 != AudioMimeType.WAV # type: ignore[comparison-overlap] + + def test_mime_type_case_sensitivity(self) -> None: + """MIME types should be case-sensitive (per RFC standard).""" + assert ImageMimeType.PNG != "IMAGE/PNG" # type: ignore[comparison-overlap] + assert ImageMimeType.PNG == "image/png" # type: ignore[comparison-overlap] + + def test_cannot_create_invalid_mime_type(self) -> None: + """Cannot instantiate invalid MIME type enum values.""" + with pytest.raises(ValueError): + ImageMimeType("invalid/type") + + def test_mime_type_hashable(self) -> None: + """MIME types should be hashable for use in sets/dicts.""" + # Important for caching and deduplication + mime_set = {ImageMimeType.PNG, ImageMimeType.JPEG, ImageMimeType.PNG} + assert len(mime_set) == 2 # PNG should be deduplicated + + mime_dict = { + ImageMimeType.PNG: "png_handler", + ImageMimeType.JPEG: "jpeg_handler", + } + assert mime_dict[ImageMimeType.PNG] == "png_handler" diff --git a/tests/unit_tests/test_models.py b/tests/unit_tests/test_models.py new file mode 100644 index 00000000..a6ede149 --- /dev/null +++ b/tests/unit_tests/test_models.py @@ -0,0 +1,338 @@ +"""Tests for models and model registry.""" + +from collections.abc import Generator +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from celeste import Capability, Model, Provider +from celeste.constraints import Str +from celeste.models import clear, get_model, list_models, register_models + +# Test data constants +SAMPLE_MODELS = [ + Model( + id="gpt-4", + provider=Provider.OPENAI, + capabilities={Capability.TEXT_GENERATION}, + display_name="GPT-4", + ), + Model( + id="dall-e-3", + provider=Provider.OPENAI, + capabilities={Capability.IMAGE_GENERATION}, + display_name="DALL-E 3", + ), + Model( + id="claude-3", + provider=Provider.ANTHROPIC, + capabilities={Capability.TEXT_GENERATION}, + display_name="Claude 3", + ), +] + + +@pytest.fixture(autouse=True) +def clear_registry() -> Generator[None, None, None]: + """Clear registry before and after each test for isolation.""" + clear() + yield + clear() + + +class TestRegisterModels: + """Test model registration functionality.""" + + @pytest.mark.smoke + def test_register_models_accepts_single_or_list(self) -> None: + """Registering models works with both single model and list.""" + single_model = SAMPLE_MODELS[0] + register_models(single_model) + assert get_model(single_model.id, single_model.provider) == single_model + + clear() + + register_models(SAMPLE_MODELS) + assert len(list_models()) == 3 + for model in SAMPLE_MODELS: + assert get_model(model.id, model.provider) == model + + def test_reregistering_same_key_raises_error(self) -> None: + """Re-registering with same (id, provider) raises ValueError.""" + original = SAMPLE_MODELS[0] + register_models(original) + + duplicate = Model( + id=original.id, + provider=original.provider, + capabilities={Capability.IMAGE_GENERATION}, + display_name="Duplicate GPT-4", + ) + + with pytest.raises(ValueError, match="already registered"): + register_models(duplicate) + + result = get_model(original.id, original.provider) + assert result == original + assert len(list_models()) == 1 + + +class TestListModels: + """Test model listing and filtering functionality.""" + + @pytest.fixture(autouse=True) + def setup_models(self) -> None: + """Set up test models for filtering tests.""" + register_models(SAMPLE_MODELS) + + def test_list_all_models(self) -> None: + """Listing all models without filters.""" + models = list_models() + assert len(models) == 3 + assert set(m.id for m in models) == {"gpt-4", "dall-e-3", "claude-3"} + + def test_filter_by_provider(self) -> None: + """Filtering models by provider.""" + openai_models = list_models(provider=Provider.OPENAI) + assert len(openai_models) == 2 + assert all(m.provider == Provider.OPENAI for m in openai_models) + + anthropic_models = list_models(provider=Provider.ANTHROPIC) + assert len(anthropic_models) == 1 + assert anthropic_models[0].id == "claude-3" + + def test_filter_by_capability(self) -> None: + """Filtering models by capability.""" + text_models = list_models(capability=Capability.TEXT_GENERATION) + assert len(text_models) == 2 + assert all(Capability.TEXT_GENERATION in m.capabilities for m in text_models) + + image_models = list_models(capability=Capability.IMAGE_GENERATION) + assert len(image_models) == 1 + assert image_models[0].id == "dall-e-3" + + def test_filter_by_both_provider_and_capability(self) -> None: + """Filtering with multiple criteria.""" + models = list_models( + provider=Provider.OPENAI, + capability=Capability.TEXT_GENERATION, + ) + assert len(models) == 1 + assert models[0].id == "gpt-4" + + @pytest.mark.parametrize( + "provider,capability", + [ + (Provider.GOOGLE, None), + (Provider.ANTHROPIC, Capability.IMAGE_GENERATION), + ], + ids=["wrong_provider", "wrong_provider_and_capability"], + ) + def test_filter_returns_empty_when_no_match( + self, provider: Provider, capability: Capability | None + ) -> None: + """Filters return empty list when no models match.""" + assert list_models(provider=provider, capability=capability) == [] + + +class TestGetModel: + """Test individual model retrieval.""" + + def test_get_existing_model(self) -> None: + """Retrieving an existing model by id and provider.""" + model = SAMPLE_MODELS[0] + register_models(model) + + result = get_model(model.id, model.provider) + assert result == model + + def test_get_nonexistent_model_from_empty_registry_returns_none(self) -> None: + """Getting a model from empty registry returns None.""" + assert get_model("nonexistent", Provider.OPENAI) is None + + @pytest.mark.parametrize( + "model_id,provider", + [("gpt-5", Provider.OPENAI), ("gpt-4", Provider.ANTHROPIC)], + ids=["wrong_id", "wrong_provider"], + ) + def test_get_model_from_populated_registry_with_wrong_key( + self, model_id: str, provider: Provider + ) -> None: + """get_model returns None for non-existent model in populated registry.""" + register_models(SAMPLE_MODELS) + assert get_model(model_id, provider) is None + + def test_same_id_different_providers_are_distinct(self) -> None: + """Models with same ID but different providers are kept distinct.""" + model1 = Model( + id="shared-id", + provider=Provider.OPENAI, + capabilities={Capability.TEXT_GENERATION}, + display_name="OpenAI Model", + ) + model2 = Model( + id="shared-id", + provider=Provider.ANTHROPIC, + capabilities={Capability.TEXT_GENERATION}, + display_name="Anthropic Model", + ) + + register_models([model1, model2]) + + assert get_model("shared-id", Provider.OPENAI) == model1 + assert get_model("shared-id", Provider.ANTHROPIC) == model2 + + +class TestEntryPoints: + """Test entry point loading functionality.""" + + @patch("celeste.importlib.metadata.entry_points") + def test_entry_point_loading_success( + self, mock_entry_points: Mock, capsys: pytest.CaptureFixture[str] + ) -> None: + """Successful loading of models from entry points.""" + mock_ep = MagicMock() + mock_ep.name = "test_models" + test_model = Model( + id="ep-test-model", + provider=Provider.OPENAI, + capabilities={Capability.TEXT_GENERATION}, + display_name="Entry Point Test Model", + ) + mock_ep.load.return_value = lambda: register_models(test_model) + + mock_entry_points.return_value = [mock_ep] + + clear() + from celeste import _load_from_entry_points + + _load_from_entry_points() + + models = list_models() + assert any(m.id == "ep-test-model" for m in models) + + captured = capsys.readouterr() + assert captured.err == "" + + @patch("celeste.importlib.metadata.entry_points") + def test_entry_point_returns_none_handled( + self, mock_entry_points: Mock, capsys: pytest.CaptureFixture[str] + ) -> None: + """Entry points returning None are handled gracefully.""" + mock_ep = MagicMock() + mock_ep.name = "empty_models" + mock_ep.load.return_value = lambda: None + + mock_entry_points.return_value = [mock_ep] + + clear() + from celeste import _load_from_entry_points + + _load_from_entry_points() + + captured = capsys.readouterr() + assert captured.err == "" + + +class TestParameterSupport: + """Test registry with models that have supported parameters.""" + + def test_register_model_with_parameters(self) -> None: + """Models with parameter_constraints are registered correctly.""" + model = Model( + id="param-model", + provider=Provider.OPENAI, + capabilities={Capability.TEXT_GENERATION}, + display_name="Model with Params", + parameter_constraints={"temperature": Str(), "max_tokens": Str()}, + ) + + register_models(model) + retrieved = get_model("param-model", Provider.OPENAI) + + assert retrieved is not None + assert len(retrieved.supported_parameters) == 2 + assert "temperature" in retrieved.supported_parameters + assert "max_tokens" in retrieved.supported_parameters + + def test_list_models_includes_parameters(self) -> None: + """list_models returns models with their parameter_constraints intact.""" + models = [ + Model( + id="with-params", + provider=Provider.ANTHROPIC, + capabilities={Capability.TEXT_GENERATION}, + display_name="With Params", + parameter_constraints={"feature_a": Str(), "feature_b": Str()}, + ), + Model( + id="without-params", + provider=Provider.GOOGLE, + capabilities={Capability.IMAGE_GENERATION}, + display_name="Without Params", + ), + ] + + register_models(models) + all_models = list_models() + + with_params = next((m for m in all_models if m.id == "with-params"), None) + without_params = next((m for m in all_models if m.id == "without-params"), None) + + assert with_params is not None + assert len(with_params.supported_parameters) == 2 + assert "feature_a" in with_params.supported_parameters + + assert without_params is not None + assert len(without_params.supported_parameters) == 0 + + +class TestClear: + """Test registry clearing functionality.""" + + def test_clear_removes_all_models(self) -> None: + """clear removes all registered models.""" + register_models(SAMPLE_MODELS) + assert len(list_models()) == 3 + + clear() + assert len(list_models()) == 0 + assert get_model("gpt-4", Provider.OPENAI) is None + + +class TestModel: + """Test Model class behavior.""" + + def test_supported_parameters_computed_from_constraints(self) -> None: + """supported_parameters property computes from parameter_constraints keys.""" + model = Model( + id="test-model", + provider=Provider.OPENAI, + capabilities={Capability.TEXT_GENERATION}, + display_name="Test Model", + parameter_constraints={"param_a": Str(), "param_b": Str()}, + ) + + register_models(model) + retrieved = get_model("test-model", Provider.OPENAI) + + assert retrieved is not None + assert len(retrieved.supported_parameters) == 2 + assert "param_a" in retrieved.supported_parameters + assert "param_b" in retrieved.supported_parameters + assert "param_c" not in retrieved.supported_parameters + + def test_supported_parameters_empty_by_default(self) -> None: + """supported_parameters defaults to empty set when parameter_constraints is empty.""" + model = Model( + id="basic-model", + provider=Provider.OPENAI, + capabilities={Capability.TEXT_GENERATION}, + display_name="Basic Model", + ) + + register_models(model) + retrieved = get_model("basic-model", Provider.OPENAI) + + assert retrieved is not None + assert retrieved.supported_parameters == set() diff --git a/tests/unit_tests/test_parameters.py b/tests/unit_tests/test_parameters.py new file mode 100644 index 00000000..de41061d --- /dev/null +++ b/tests/unit_tests/test_parameters.py @@ -0,0 +1,58 @@ +"""High-value tests for celeste.parameters module.""" + +from enum import StrEnum +from typing import Any + +import pytest + +from celeste.core import Parameter +from celeste.models import Model + + +class DefaultParseOutputMapper: + """Mapper that uses default parse_output behavior (returns content unchanged).""" + + name: StrEnum = Parameter.TEMPERATURE + + def map(self, request: dict[str, Any], value: Any, model: Model) -> dict[str, Any]: # noqa: ANN401 + """Transform parameter value into provider request.""" + if value is not None: + request["temperature"] = value + return request + + def parse_output(self, content: object, value: object | None) -> object: + """Default implementation: return content unchanged.""" + return content + + +class TestParameterMapperProtocol: + """Test ParameterMapper Protocol behavior.""" + + @pytest.mark.parametrize( + "content,value", + [ + ("test string", None), + ({"key": "value"}, None), + ([1, 2, 3], None), + (None, None), + (42, None), + (True, None), + ("test content", "some parameter value"), + ], + ) + def test_parse_output_returns_content_unchanged( + self, content: object, value: object | None + ) -> None: + """Default parse_output implementation returns content unchanged. + + Tests the documented default behavior of ParameterMapper.parse_output. + This covers line 50 in parameters.py. + """ + # Arrange + mapper = DefaultParseOutputMapper() + + # Act + result = mapper.parse_output(content, value=value) + + # Assert + assert result is content diff --git a/tests/unit_tests/test_streaming.py b/tests/unit_tests/test_streaming.py new file mode 100644 index 00000000..54ba5d1e --- /dev/null +++ b/tests/unit_tests/test_streaming.py @@ -0,0 +1,688 @@ +"""High-value tests for Stream - focusing on lifecycle, resource cleanup, and state management.""" + +from collections.abc import AsyncIterator +from typing import Any +from unittest.mock import AsyncMock + +import pytest +from pydantic import Field + +from celeste.io import Chunk, FinishReason, Output, Usage +from celeste.streaming import Stream + + +class ConcreteOutput(Output[str]): + """Concrete output for testing Stream.""" + + pass + + +class ConcreteStream(Stream[ConcreteOutput]): + """Concrete Stream implementation for testing abstract behavior.""" + + def __init__( + self, + sse_iterator: AsyncIterator[dict[str, Any]], + chunk_events: list[dict[str, Any]] | None = None, + filter_none: bool = False, + ) -> None: + """Initialize stream with configurable parsing behavior. + + Args: + sse_iterator: Server-Sent Events iterator. + chunk_events: Events that should be converted to chunks. + filter_none: If True, return None for events not in chunk_events. + """ + super().__init__(sse_iterator) + self._chunk_events = chunk_events or [] + self._filter_none = filter_none + + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: + """Parse event to Chunk or None (for lifecycle events).""" + # If filter_none enabled, only return chunks for events in chunk_events + if self._filter_none and event not in self._chunk_events: + return None + + content = event.get("delta", event.get("content", "")) + if not content: + return None + + return Chunk( + content=content, + finish_reason=None, # Base tests don't use typed finish reasons + metadata=event.get("metadata", {}), + ) + + def _parse_output(self, chunks: list[Chunk]) -> ConcreteOutput: + """Combine chunks into final output.""" + content = "".join(chunk.content for chunk in chunks) + # Usage is empty base class - store data in metadata for tests + usage = Usage() + return ConcreteOutput( + content=content, + usage=usage, + metadata=chunks[-1].metadata if chunks else {}, + ) + + +@pytest.fixture +def mock_sse_iterator() -> AsyncMock: + """Create mock SSE iterator with aclose capability.""" + mock_iter = AsyncMock(spec=["aclose", "__anext__"]) + mock_iter.aclose = AsyncMock() + return mock_iter + + +async def _async_iter(items: list[dict[str, Any]]) -> AsyncIterator[dict[str, Any]]: + """Convert list to async iterator.""" + for item in items: + yield item + + +class TestStreamAsyncIterator: + """Test Stream AsyncIterator protocol - chunk filtering and exhaustion.""" + + async def test_iteration_yields_chunks_and_filters_none(self) -> None: + """Stream must filter None chunks (lifecycle events) and only yield valid chunks.""" + # Arrange - Mix of valid chunks and lifecycle events + events = [ + {"delta": "Hello"}, # Valid chunk + {"type": "ping"}, # Lifecycle event (no delta) + {"delta": " world"}, # Valid chunk + {"finish_reason": "stop"}, # Lifecycle event (no delta) + ] + stream = ConcreteStream(_async_iter(events)) + + # Act - Collect chunks via iteration + chunks = [chunk async for chunk in stream] + + # Assert - Only valid chunks yielded, lifecycle events filtered + assert len(chunks) == 2 + assert chunks[0].content == "Hello" + assert chunks[1].content == " world" + + async def test_iteration_stops_when_sse_exhausted(self) -> None: + """Stream must raise StopAsyncIteration when SSE iterator is exhausted.""" + # Arrange + events = [{"delta": "test"}] + stream = ConcreteStream(_async_iter(events)) + + # Act - Consume all chunks + chunks = [chunk async for chunk in stream] + + # Assert - One chunk yielded + assert len(chunks) == 1 + + # Act & Assert - Next iteration raises StopAsyncIteration + with pytest.raises(StopAsyncIteration): + await stream.__anext__() + + async def test_iteration_when_closed_raises_stop_iteration(self) -> None: + """Stream must raise StopAsyncIteration immediately if already closed.""" + # Arrange + events = [{"delta": "test"}] + stream = ConcreteStream(_async_iter(events)) + await stream.aclose() + + # Act & Assert - Closed stream raises StopAsyncIteration + with pytest.raises(StopAsyncIteration): + await stream.__anext__() + + async def test_stream_accumulates_chunks_internally(self) -> None: + """Stream must accumulate all non-None chunks for later output parsing.""" + # Arrange + events = [ + {"delta": "First"}, + {"type": "lifecycle"}, # Filtered + {"delta": "Second"}, + {"delta": "Third"}, + ] + stream = ConcreteStream(_async_iter(events)) + + # Act - Consume stream + async for _ in stream: + pass + + # Assert - Verify internal accumulation via output + assert stream.output.content == "FirstSecondThird" + + +class TestStreamOutputProperty: + """Test Stream.output property - access guard and final result.""" + + async def test_output_access_before_exhaustion_raises_error(self) -> None: + """Accessing .output before stream exhaustion must raise RuntimeError.""" + # Arrange + events = [{"delta": "test"}] + stream = ConcreteStream(_async_iter(events)) + + # Act & Assert - Premature access raises RuntimeError + with pytest.raises( + RuntimeError, match=r"Stream not exhausted\. Consume all chunks" + ): + _ = stream.output + + async def test_output_accessible_after_exhaustion(self) -> None: + """Stream.output must be accessible after stream is fully consumed.""" + # Arrange + events = [ + {"delta": "Hello", "metadata": {"usage": {"input_tokens": 5}}}, + { + "delta": " world", + "finish_reason": "stop", + "metadata": {"usage": {"output_tokens": 10}}, + }, + ] + stream = ConcreteStream(_async_iter(events)) + + # Act - Consume stream + async for _ in stream: + pass + + # Assert - Output accessible with accumulated content + output = stream.output + assert output.content == "Hello world" + # Usage is empty base class - data stored in metadata + assert output.metadata.get("usage", {}).get("output_tokens") == 10 + + async def test_output_returns_consistent_result_after_exhaustion(self) -> None: + """Stream.output must return the same result on multiple accesses.""" + # Arrange + events = [{"delta": "test"}] + stream = ConcreteStream(_async_iter(events)) + + # Act - Exhaust stream + async for _ in stream: + pass + + # Assert - Multiple accesses return identical result + output1 = stream.output + output2 = stream.output + assert output1 is output2 + assert output1.content == "test" + + +class TestStreamResourceCleanup: + """Test Stream resource cleanup - aclose() idempotence and SSE cleanup.""" + + async def test_aclose_is_idempotent(self, mock_sse_iterator: AsyncMock) -> None: + """Stream.aclose() must be safe to call multiple times without error.""" + # Arrange + stream = ConcreteStream(mock_sse_iterator) + + # Act - Call aclose multiple times + await stream.aclose() + await stream.aclose() + await stream.aclose() + + # Assert - No errors raised, SSE iterator closed once + assert stream._closed is True + mock_sse_iterator.aclose.assert_called_once() + + async def test_aclose_calls_sse_iterator_aclose( + self, mock_sse_iterator: AsyncMock + ) -> None: + """Stream.aclose() must call aclose on SSE iterator if available.""" + # Arrange + stream = ConcreteStream(mock_sse_iterator) + + # Act + await stream.aclose() + + # Assert - SSE iterator cleanup called + mock_sse_iterator.aclose.assert_called_once() + + async def test_aclose_handles_sse_iterator_without_aclose(self) -> None: + """Stream.aclose() must handle SSE iterators without aclose method.""" + + # Arrange - Iterator without aclose + async def simple_iter() -> AsyncIterator[dict[str, Any]]: + yield {"delta": "test"} + + stream = ConcreteStream(simple_iter()) + + # Act & Assert - Should not raise AttributeError + await stream.aclose() + assert stream._closed is True + + async def test_aclose_called_on_natural_exhaustion(self) -> None: + """Stream must call aclose automatically when naturally exhausted.""" + # Arrange + events = [{"delta": "test"}] + stream = ConcreteStream(_async_iter(events)) + + # Act - Consume stream naturally + async for _ in stream: + pass + + # Assert - Stream marked as closed after exhaustion + assert stream._closed is True + + +class TestStreamAsyncContextManager: + """Test Stream AsyncContextManager protocol - cleanup guarantee and exception handling.""" + + async def test_context_manager_returns_self(self) -> None: + """__aenter__ must return Stream instance for iteration.""" + # Arrange + events = [{"delta": "test"}] + stream = ConcreteStream(_async_iter(events)) + + # Act & Assert + async with stream as ctx_stream: + assert ctx_stream is stream + + async def test_context_manager_calls_aclose_on_exit( + self, mock_sse_iterator: AsyncMock + ) -> None: + """__aexit__ must call aclose() to cleanup resources on exit.""" + # Arrange + stream = ConcreteStream(mock_sse_iterator) + + # Act - Enter and exit context + async with stream: + pass + + # Assert - Cleanup called + assert stream._closed is True + mock_sse_iterator.aclose.assert_called_once() + + async def test_context_manager_cleanup_on_exception( + self, mock_sse_iterator: AsyncMock + ) -> None: + """__aexit__ must cleanup resources even when exception occurs.""" + # Arrange + stream = ConcreteStream(mock_sse_iterator) + + # Act & Assert - Exception propagates but cleanup happens + with pytest.raises(ValueError, match="Test error"): + async with stream: + raise ValueError("Test error") + + # Assert - Cleanup happened despite exception + assert stream._closed is True + mock_sse_iterator.aclose.assert_called_once() + + async def test_context_manager_propagates_exceptions(self) -> None: + """__aexit__ must propagate exceptions (returns False).""" + # Arrange + events = [{"delta": "test"}] + stream = ConcreteStream(_async_iter(events)) + exc_type = RuntimeError + exc_val = RuntimeError("Intentional error") + + # Act & Assert - Exception raised in context must propagate + with pytest.raises(RuntimeError, match="Intentional error"): + async with stream: + raise exc_val + + # Assert - __aexit__ returns False to propagate exceptions + result = await stream.__aexit__(exc_type, exc_val, None) + assert result is False + + async def test_context_manager_guarantees_cleanup_on_early_break(self) -> None: + """Context manager must cleanup when iteration breaks early.""" + # Arrange + events = [ + {"delta": "first"}, + {"delta": "second"}, + {"delta": "third"}, + ] + stream = ConcreteStream(_async_iter(events)) + + # Act - Break iteration early + async with stream: + async for chunk in stream: + if chunk.content == "first": + break # Early exit + + # Assert - Cleanup still happened + assert stream._closed is True + + +class TestStreamAbstractBehavior: + """Test Stream abstract method enforcement.""" + + async def test_cannot_instantiate_base_stream_directly( + self, mock_sse_iterator: AsyncMock + ) -> None: + """Stream must not be instantiable due to abstract methods.""" + # Act & Assert + with pytest.raises( + TypeError, match=r"Can't instantiate abstract class.*Stream" + ) as exc_info: + Stream(mock_sse_iterator) # type: ignore[abstract] + + # Verify error mentions missing abstract methods + error_msg = str(exc_info.value) + assert "Stream" in error_msg + assert any(method in error_msg for method in ["_parse_chunk", "_parse_output"]) + + async def test_subclass_without_parse_chunk_fails( + self, mock_sse_iterator: AsyncMock + ) -> None: + """Subclass without _parse_chunk implementation must fail instantiation.""" + + # Arrange + class IncompleteStream(Stream[ConcreteOutput]): + """Missing _parse_chunk implementation.""" + + def _parse_output(self, chunks: list[Chunk]) -> ConcreteOutput: + """Implement _parse_output but not _parse_chunk.""" + return ConcreteOutput(content="test") + + # Act & Assert + with pytest.raises(TypeError, match=r".*_parse_chunk.*"): + IncompleteStream(mock_sse_iterator) # type: ignore[abstract] + + async def test_subclass_without_parse_output_fails( + self, mock_sse_iterator: AsyncMock + ) -> None: + """Subclass without _parse_output implementation must fail instantiation.""" + + # Arrange + class IncompleteStream(Stream[ConcreteOutput]): + """Missing _parse_output implementation.""" + + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: + """Implement _parse_chunk but not _parse_output.""" + return Chunk(content="test") + + # Act & Assert + with pytest.raises(TypeError, match=r".*_parse_output.*"): + IncompleteStream(mock_sse_iterator) # type: ignore[abstract] + + +class TestStreamRepr: + """Test Stream.__repr__ method - all state branches.""" + + async def test_repr_idle_state(self) -> None: + """__repr__ must show 'idle' state when stream initialized but no chunks yet.""" + # Arrange + events: list[dict[str, Any]] = [] + stream = ConcreteStream(_async_iter(events)) + + # Act + repr_str = repr(stream) + + # Assert + assert "idle" in repr_str + assert "chunks" not in repr_str + assert stream.__class__.__name__ in repr_str + + @pytest.mark.parametrize( + "chunk_count,expected_count", + [ + (1, "1 chunks"), + (2, "2 chunks"), + (3, "3 chunks"), + ], + ) + async def test_repr_streaming_state_shows_chunk_count( + self, chunk_count: int, expected_count: str + ) -> None: + """__repr__ must show 'streaming' state with correct chunk count.""" + # Arrange + events = [{"delta": f"chunk{i}"} for i in range(chunk_count)] + stream = ConcreteStream(_async_iter(events)) + + # Act - Get chunks but don't exhaust stream + async for _chunk in stream: + if len(stream._chunks) == chunk_count: + break # Stop before exhaustion + + # Assert - State is streaming with correct chunk count + repr_str = repr(stream) + assert "streaming" in repr_str + assert expected_count in repr_str + assert stream.__class__.__name__ in repr_str + + async def test_repr_closed_state(self) -> None: + """__repr__ must show 'closed' state when stream closed but no output set.""" + # Arrange + events = [{"delta": "test"}] + stream = ConcreteStream(_async_iter(events)) + + # Act - Close stream before exhaustion + await stream.aclose() + + # Assert - State is closed + repr_str = repr(stream) + assert "closed" in repr_str + assert stream.__class__.__name__ in repr_str + + async def test_repr_done_state(self) -> None: + """__repr__ must show 'done' state when stream exhausted with output set.""" + # Arrange + events = [{"delta": "test"}] + stream = ConcreteStream(_async_iter(events)) + + # Act - Exhaust stream + async for _ in stream: + pass + + # Assert - State is done with chunk count + repr_str = repr(stream) + assert "done" in repr_str + assert "1 chunks" in repr_str + assert stream.__class__.__name__ in repr_str + + +class TestStreamEmptyStreamError: + """Test Stream empty stream error handling.""" + + async def test_empty_stream_raises_runtime_error(self) -> None: + """Stream must raise RuntimeError when exhausted with no chunks.""" + + # Arrange - Create stream where all events return None from _parse_chunk + async def empty_iter() -> AsyncIterator[dict[str, Any]]: + yield {"type": "ping"} # Lifecycle event (no delta) + yield {"finish_reason": "stop"} # Lifecycle event (no delta) + + stream = ConcreteStream(empty_iter()) + + # Act & Assert - Exhaustion raises RuntimeError + with pytest.raises( + RuntimeError, match=r"Stream completed but no chunks were produced" + ): + async for _ in stream: + pass + + async def test_stream_with_only_lifecycle_events_raises_error(self) -> None: + """Stream raises RuntimeError when SSE yields events but all chunks are filtered to None.""" + # Arrange - Events that all return None from _parse_chunk + events = [ + {"type": "ping"}, # Lifecycle event (no delta/content) + {"type": "start"}, # Lifecycle event + {"type": "end"}, # Lifecycle event + {"finish_reason": "stop"}, # Lifecycle event + ] + stream = ConcreteStream(_async_iter(events)) + + # Act & Assert - Should raise RuntimeError when exhausted + with pytest.raises( + RuntimeError, match=r"Stream completed but no chunks were produced" + ): + async for _ in stream: + pass + + +class TestStreamExceptionHandling: + """Test Stream exception handling in __anext__.""" + + async def test_anext_handles_exception_in_parse_chunk_and_cleans_up(self) -> None: + """__anext__ must call aclose() and re-raise when _parse_chunk raises exception.""" + + # Arrange - Stream that raises exception in _parse_chunk + class ExceptionStream(ConcreteStream): + """Stream that raises exception in _parse_chunk.""" + + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: + """Raise exception during parsing.""" + raise ValueError("Parse error") + + events = [{"delta": "test"}] + stream = ExceptionStream(_async_iter(events)) + + # Act & Assert - Exception is re-raised and cleanup happens + with pytest.raises(ValueError, match="Parse error"): + await stream.__anext__() + + # Assert - Cleanup was called + assert stream._closed is True + + async def test_anext_handles_exception_in_parse_output_and_cleans_up(self) -> None: + """__anext__ must call aclose() and re-raise when _parse_output raises exception.""" + + # Arrange - Stream that raises exception in _parse_output + class ExceptionStream(ConcreteStream): + """Stream that raises exception in _parse_output.""" + + def _parse_output(self, chunks: list[Chunk]) -> ConcreteOutput: + """Raise exception during output parsing.""" + raise RuntimeError("Output parse error") + + events = [{"delta": "test"}] + stream = ExceptionStream(_async_iter(events)) + + # Act & Assert - Exception is re-raised and cleanup happens + with pytest.raises(RuntimeError, match="Output parse error"): + async for _ in stream: + pass + + # Assert - Cleanup was called + assert stream._closed is True + + async def test_anext_handles_exception_in_sse_iterator_and_cleans_up(self) -> None: + """__anext__ must call aclose() and re-raise when SSE iterator raises exception.""" + + # Arrange - SSE iterator that raises exception + async def failing_iter() -> AsyncIterator[dict[str, Any]]: + """Iterator that raises exception.""" + yield {"delta": "test"} + raise ConnectionError("Connection lost") + + stream = ConcreteStream(failing_iter()) + + # Act & Assert - Exception is re-raised and cleanup happens + with pytest.raises(ConnectionError, match="Connection lost"): + async for chunk in stream: + # First chunk succeeds + assert chunk.content == "test" + # Next iteration will raise exception + await stream.__anext__() + + # Assert - Cleanup was called + assert stream._closed is True + + +class TestStreamWithTypedUsageAndFinishReason: + """Test Stream with capability-specific Usage and FinishReason subclasses. + + Validates that the hierarchical pattern works with actual fields + before migrating existing packages (CEL-39). + """ + + async def test_typed_usage_and_finish_reason_in_chunks(self) -> None: + """Stream must support capability-specific Usage/FinishReason with actual fields.""" + + # Arrange - Define typed Usage and FinishReason with fields + class TypedUsage(Usage): + """Typed usage with actual fields.""" + + input_tokens: int + output_tokens: int + + class TypedFinishReason(FinishReason): + """Typed finish reason with actual field.""" + + reason: str + + class TypedChunk(Chunk[str]): + """Typed chunk with specific finish_reason and usage.""" + + finish_reason: TypedFinishReason | None = None + usage: TypedUsage | None = None + + class TypedOutput(Output[str]): + """Typed output with specific usage.""" + + usage: TypedUsage = Field( + default_factory=lambda: TypedUsage(input_tokens=0, output_tokens=0) + ) + + class TypedStream(Stream[TypedOutput]): + """Stream using typed classes.""" + + def _parse_chunk(self, event: dict[str, Any]) -> TypedChunk | None: + """Parse event to typed chunk.""" + content = event.get("delta") + if not content: + return None + + # Build typed finish_reason if present + finish_reason = None + if "finish_reason" in event: + finish_reason = TypedFinishReason(reason=event["finish_reason"]) + + # Build typed usage if present + usage = None + if "usage" in event: + usage = TypedUsage( + input_tokens=event["usage"]["input_tokens"], + output_tokens=event["usage"]["output_tokens"], + ) + + return TypedChunk( + content=content, + finish_reason=finish_reason, + usage=usage, + ) + + def _parse_output(self, chunks: list[Chunk]) -> TypedOutput: + """Combine chunks into typed output.""" + content = "".join(chunk.content for chunk in chunks) + + # Extract usage from final chunk (typed) + final_chunk = chunks[-1] if chunks else None + usage = ( + final_chunk.usage + if final_chunk and isinstance(final_chunk.usage, TypedUsage) + else TypedUsage(input_tokens=0, output_tokens=0) + ) + + return TypedOutput(content=content, usage=usage) + + # Act - Stream with typed chunks + events: list[dict[str, Any]] = [ + {"delta": "Hello"}, + {"delta": " world"}, + { + "delta": "!", + "finish_reason": "stop", + "usage": {"input_tokens": 5, "output_tokens": 10}, + }, + ] + stream = TypedStream(_async_iter(events)) + + chunks = [chunk async for chunk in stream] + + # Assert - Chunks have typed finish_reason and usage + assert len(chunks) == 3 + assert chunks[0].finish_reason is None + assert chunks[0].usage is None + + assert chunks[2].finish_reason is not None + assert isinstance(chunks[2].finish_reason, TypedFinishReason) + assert chunks[2].finish_reason.reason == "stop" + assert chunks[2].usage is not None + assert isinstance(chunks[2].usage, TypedUsage) + assert chunks[2].usage.input_tokens == 5 + assert chunks[2].usage.output_tokens == 10 + + # Assert - Output has typed usage + output = stream.output + assert output.content == "Hello world!" + assert isinstance(output.usage, TypedUsage) + assert output.usage.input_tokens == 5 + assert output.usage.output_tokens == 10