| from pathlib import Path |
|
|
| from aistudio_api.infrastructure.gateway.wire_codec import AistudioWireCodec |
| from aistudio_api.infrastructure.gateway.wire_types import ( |
| AistudioGenerationConfig, |
| AistudioImageOutputMode, |
| AistudioThinkingConfig, |
| MediaResolution, |
| ThinkingLevel, |
| ) |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def test_decode_image_request_maps_generation_config_fields_from_proto_indexes(): |
| codec = AistudioWireCodec() |
| raw = (ROOT / "test-image-input.json").read_text() |
|
|
| request = codec.decode(raw) |
|
|
| assert request.model == "models/gemini-3.1-flash-image-preview" |
| assert request.generation_config.stop_sequences == ["6"] |
| assert request.generation_config.max_tokens == 65536 |
| assert request.generation_config.temperature == 1 |
| assert request.generation_config.top_p == 0.95 |
| assert request.generation_config.top_k == 64 |
| assert request.generation_config.image_output_mode == [2, 1] |
| assert request.generation_config.thinking_config == [1, None, None, 3] |
| assert request.request_flag == 1 |
| assert request.cached_content.startswith("v1_") |
|
|
|
|
| def test_encode_preserves_newly_mapped_proto_fields(): |
| codec = AistudioWireCodec() |
| raw = (ROOT / "test-image-input.json").read_text() |
|
|
| request = codec.decode(raw) |
| encoded = codec.encode(request) |
| reparsed = codec.decode(encoded) |
|
|
| assert reparsed.generation_config.stop_sequences == ["6"] |
| assert reparsed.generation_config.image_output_mode == [2, 1] |
| assert reparsed.generation_config.thinking_config == [1, None, None, 3] |
| assert reparsed.request_flag == 1 |
| assert reparsed.cached_content == request.cached_content |
|
|
|
|
| def test_thinking_config_encodes_high_level_wire_shape(): |
| assert AistudioThinkingConfig(ThinkingLevel.HIGH).to_wire() == [1, None, None, 3] |
|
|
|
|
| def test_generation_config_enables_default_thinking(): |
| config = AistudioGenerationConfig([]) |
|
|
| config.enable_default_thinking() |
|
|
| assert config.thinking_config == [1, None, None, 3] |
| assert config.media_resolution is None |
|
|
|
|
| def test_generation_config_accepts_readable_image_output_mode_wrapper(): |
| config = AistudioGenerationConfig([]) |
|
|
| config.image_output_mode = AistudioImageOutputMode.text_and_image() |
|
|
| assert config.image_output_mode == [2, 1] |
|
|
|
|
| def test_generation_config_accepts_media_resolution_enum(): |
| config = AistudioGenerationConfig([]) |
|
|
| config.media_resolution = MediaResolution.MEDIUM |
|
|
| assert config.media_resolution == 2 |
|
|
