File size: 2,485 Bytes
95f8c4b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | 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
|