import base64 import hashlib import json import threading from io import BytesIO from types import SimpleNamespace import google.auth import pytest import torch from PIL import Image from custom_nodes.ComfyUI_Gemini_Expanded_API import interactions_nodes from custom_nodes.ComfyUI_Gemini_Expanded_API.response_schema import SSL_GeminiResponseSchema, default_state, default_entry @pytest.fixture def mock_interactions_client(monkeypatch): node = interactions_nodes.SSL_GeminiInteractionsTextPrompt for cache in (node._cache, node._client_cache, node._seed_map_cache): cache.clear() captured = SimpleNamespace(calls=[], clients=[], responses=[], on_client=None) class FakeInteractions: def create(self, **kwargs): captured.calls.append(kwargs) resp = captured.responses.pop(0) if captured.responses else {"text": "default response", "id": "int_default_123"} if callable(resp): resp = resp() if isinstance(resp, Exception): raise resp steps = [] if "thoughts" in resp: steps.append(SimpleNamespace( type="thought", summary=[SimpleNamespace(type="text", text=resp["thoughts"])], text="", )) if "steps" in resp: steps.extend(resp["steps"]) elif "text" in resp: steps.append(SimpleNamespace( type="model_output", content=[SimpleNamespace(type="text", text=resp["text"])], )) out_img = None if "image_data" in resp: out_img = SimpleNamespace( type="image", data=resp["image_data"], mime_type=resp.get("image_mime", "image/png"), ) return SimpleNamespace( id=resp.get("id", "int_123"), status=resp.get("status", "completed"), errors=resp.get("errors", None), output_text=resp.get("text", ""), output_image=out_img, steps=steps, usage=SimpleNamespace(total_cached_tokens=resp.get("cached_tokens", 0)) if "cached_tokens" in resp else None, ) class FakeClient: def __init__(self, **kwargs): captured.clients.append(kwargs) if captured.on_client: captured.on_client() self.interactions = FakeInteractions() monkeypatch.setattr(interactions_nodes.genai, "Client", FakeClient) yield captured for cache in (node._cache, node._client_cache, node._seed_map_cache): cache.clear() def _config(**overrides): config = { "api_key": "test_api_key", "api_version": "v1beta", "use_vertexai_env": False, "vertexai_express": False, "vertexai_project": "", "vertexai_location": "", "google_application_credentials": "", } config.update(overrides) return config def _execute_kwargs(config=None, **overrides): kwargs = { "config": config or _config(), "prompt": "test prompt", "system_instruction": "test system instruction", "model": "gemini-2.5-flash", "temperature": 1.0, "top_p": 0.95, "top_k": 40, "max_output_tokens": 128, "include_images": False, "aspect_ratio": "None", "bypass_mode": "None", "thinking_budget": 0, "use_proxy": False, "proxy_host": "127.0.0.1", "proxy_port": 7890, "use_seed": False, "seed": 0, "timeout": 30, "include_thoughts": False, "thinking_level": "None", "media_resolution": "unspecified", "retry_pattern": "", "max_retries": 3, "timeout_fallback_text": "", "image_size": "None", "store": True, "previous_interaction_id": "", "video": None, "image_inputs": None, "response_schema": None, } kwargs.update(overrides) return kwargs class FakeVideo: def __init__(self, data): self.data = data self.save_calls = [] def save_to(self, buffer, format, codec): self.save_calls.append((format, codec)) buffer.write(self.data) def test_interactions_define_schema(): schema = interactions_nodes.SSL_GeminiInteractionsTextPrompt.define_schema() assert schema.node_id == "SSL_GeminiInteractionsTextPrompt" assert schema.category == "API/Gemini/Interactions" input_names = [i.id for i in schema.inputs] assert "store" in input_names assert "previous_interaction_id" in input_names assert "response_schema" in input_names assert "image_inputs" in input_names output_names = [o.id for o in schema.outputs] assert output_names == ["text", "image", "final_actual_seed", "structured_output", "thoughts", "interaction_id"] def test_interactions_api_key_config_node(): schema = interactions_nodes.SSL_GeminiInteractionsAPIKeyConfig.define_schema() assert schema.node_id == "SSL_GeminiInteractionsAPIKeyConfig" assert schema.category == "API/Gemini/Interactions" output = interactions_nodes.SSL_GeminiInteractionsAPIKeyConfig.execute( api_key="secret123", api_version="v1beta", use_vertexai_env=False, vertexai_express=True, ) config = output[0] assert config["api_key"] == "secret123" assert config["api_version"] == "v1beta" assert config["vertexai_express"] is True def test_basic_text_generation(mock_interactions_client): mock_interactions_client.responses = [{"text": "Hello world!", "id": "int_abc_123"}] args = _execute_kwargs() output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert output[0] == "Hello world!" # text assert output[5] == "int_abc_123" # interaction_id assert len(mock_interactions_client.calls) == 1 call = mock_interactions_client.calls[0] assert call["model"] == "gemini-2.5-flash" assert call["input"] == "test prompt" assert call["system_instruction"] == "test system instruction" assert call["store"] is True assert "safety_settings" in call assert len(call["safety_settings"]) == 10 assert {s["type"] for s in call["safety_settings"]} == set(interactions_nodes.SSL_GeminiInteractionsTextPrompt.ALL_HARM_CATEGORIES) assert all(s["threshold"] == "block_none" for s in call["safety_settings"]) def test_multimodal_image_input(mock_interactions_client): mock_interactions_client.responses = [{"text": "Found a square in image", "id": "int_img_1"}] image_tensor = torch.zeros((1, 32, 32, 3), dtype=torch.float32) args = _execute_kwargs( image_inputs={"image_1": image_tensor}, prompt="Describe this image", media_resolution="high", ) output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert output[0] == "Found a square in image" call = mock_interactions_client.calls[0] assert isinstance(call["input"], list) assert len(call["input"]) == 2 img_item = call["input"][0] assert img_item["type"] == "image" assert img_item["mime_type"] == "image/png" assert isinstance(img_item["data"], BytesIO) assert img_item["data"].tell() == 0 assert len(img_item["data"].getvalue()) > 0 assert img_item["resolution"] == "high" from google.genai._gaos.types.interactions.imagecontent import ImageContent validated_img = ImageContent(**img_item) assert isinstance(validated_img.data, str) assert len(validated_img.data) > 0 text_item = call["input"][1] assert text_item["type"] == "text" assert text_item["text"] == "Describe this image" def test_multimodal_video_input(mock_interactions_client): mock_interactions_client.responses = [{"text": "Found motion in video", "id": "int_vid_1"}] video_data = b"fake mp4 video bytes" args = _execute_kwargs( video={"video": FakeVideo(video_data), "fps": 2, "pad_at_start": False, "duration_aware_padding": False}, prompt="What happened in the video?", ) output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert output[0] == "Found motion in video" call = mock_interactions_client.calls[0] assert isinstance(call["input"], list) assert len(call["input"]) == 2 vid_item = call["input"][0] assert vid_item["type"] == "video" assert vid_item["mime_type"] == "video/mp4" assert isinstance(vid_item["data"], BytesIO) assert vid_item["data"].getvalue() == video_data from google.genai._gaos.types.interactions.videocontent import VideoContent validated_vid = VideoContent(**vid_item) assert base64.b64decode(validated_vid.data) == video_data def test_multimodal_image_and_video_ordering(mock_interactions_client): mock_interactions_client.responses = [{"text": "Found image and video", "id": "int_both_1"}] video_data = b"fake video bytes" image_tensor = torch.zeros(1, 16, 16, 3) args = _execute_kwargs( video={"video": FakeVideo(video_data), "fps": 2, "pad_at_start": False, "duration_aware_padding": False}, image_inputs={"image_1": image_tensor}, prompt="Describe image and video", ) output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert output[0] == "Found image and video" call = mock_interactions_client.calls[0] items = call["input"] assert len(items) == 3 assert items[0]["type"] == "image" assert items[1]["type"] == "video" assert items[2]["type"] == "text" def test_structured_output_with_schema(mock_interactions_client): schema = { "type": "object", "properties": {"animal": {"type": "string"}, "count": {"type": "integer"}}, "required": ["animal", "count"], } mock_interactions_client.responses = [{"text": '{"animal": "cat", "count": 3}', "id": "int_struct_1"}] args = _execute_kwargs(response_schema=schema) output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert output[0] == '{"animal": "cat", "count": 3}' assert output[3] == '{"animal": "cat", "count": 3}' # structured_output call = mock_interactions_client.calls[0] assert "response_format" in call assert call["response_format"] == [ {"type": "text", "mime_type": "application/json", "schema": schema} ] def test_image_generation(mock_interactions_client): # Create small valid 16x16 PNG image img = Image.new("RGB", (16, 16), color=(255, 0, 0)) buf = BytesIO() img.save(buf, format="PNG") b64_png = base64.b64encode(buf.getvalue()).decode("utf-8") mock_interactions_client.responses = [{ "text": "Generated an image", "image_data": b64_png, "image_mime": "image/png", "id": "int_img_gen_1", }] args = _execute_kwargs( model="gemini-3.1-flash-image", include_images=True, aspect_ratio="16:9", image_size="1K", ) output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) image_tensor = output[1] assert torch.is_tensor(image_tensor) assert image_tensor.shape == (1, 16, 16, 3) call = mock_interactions_client.calls[0] assert call["response_format"] == [ {"type": "image", "aspect_ratio": "16:9", "image_size": "1K"} ] def test_thinking_and_thoughts_extraction(mock_interactions_client): mock_interactions_client.responses = [{ "text": "42 is the answer", "thoughts": "I need to calculate the answer carefully.", "id": "int_think_1", }] args = _execute_kwargs( model="gemini-3.8-flash", thinking_level="high", include_thoughts=True, ) output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert output[0] == "42 is the answer" assert output[4] == "I need to calculate the answer carefully." call = mock_interactions_client.calls[0] assert call["generation_config"]["thinking_level"] == "high" assert call["generation_config"]["thinking_summaries"] == "auto" def test_multi_turn_chaining(mock_interactions_client): mock_interactions_client.responses = [ {"text": "Turn 1 answer", "id": "int_turn_1"}, {"text": "Turn 2 answer", "id": "int_turn_2"}, ] # First turn args1 = _execute_kwargs(prompt="First message", store=True) out1 = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args1) id1 = out1[5] assert id1 == "int_turn_1" assert mock_interactions_client.calls[0].get("previous_interaction_id") is None # Second turn continuing first args2 = _execute_kwargs(prompt="Follow-up message", previous_interaction_id=id1, store=True) out2 = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args2) assert out2[0] == "Turn 2 answer" assert out2[5] == "int_turn_2" assert mock_interactions_client.calls[1]["previous_interaction_id"] == "int_turn_1" def test_vertexai_enterprise_auth(mock_interactions_client, monkeypatch): monkeypatch.setenv("GOOGLE_GENAI_USE_ENTERPRISE", "true") monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "test-project") monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "us-central1") config = _config(use_vertexai_env=True) args = _execute_kwargs(config=config) interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) client_args = mock_interactions_client.clients[0] assert client_args.get("enterprise") is True assert client_args.get("project") == "test-project" assert client_args.get("location") == "us-central1" def test_vertexai_express_auth(mock_interactions_client): config = _config(vertexai_express=True, api_key="my_express_key") args = _execute_kwargs(config=config) interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) client_args = mock_interactions_client.clients[0] assert client_args.get("enterprise") is True assert client_args.get("api_key") == "my_express_key" def test_local_cache_hit_on_same_seed(mock_interactions_client): mock_interactions_client.responses = [ {"text": "Cached response", "id": "int_cache_1"}, ] args = _execute_kwargs(use_seed=True, seed=42) out1 = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert out1[0] == "Cached response" assert len(mock_interactions_client.calls) == 1 # Second call with same seed and params should hit local cache out2 = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert out2[0] == "Cached response" assert len(mock_interactions_client.calls) == 1 def test_retry_pattern_triggers_retry(mock_interactions_client): mock_interactions_client.responses = [ {"text": "Please retry me", "id": "int_r1"}, {"text": "Final good answer", "id": "int_r2"}, ] args = _execute_kwargs( retry_pattern="retry me", max_retries=2, use_seed=True, seed=100, ) output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert output[0] == "Final good answer" assert len(mock_interactions_client.calls) == 2 def test_timeout_fallback(mock_interactions_client, monkeypatch): def slow_create(**kwargs): import time time.sleep(1.0) return SimpleNamespace(id="slow", output_text="slow") class SlowInteractions: def create(self, **kwargs): return slow_create(**kwargs) class SlowClient: def __init__(self, **kwargs): self.interactions = SlowInteractions() monkeypatch.setattr(interactions_nodes.genai, "Client", SlowClient) args = _execute_kwargs(timeout=0.1, timeout_fallback_text="Custom timeout message") output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert output[0] == "Custom timeout message" def test_interactions_api_key_config_with_cache(): output = interactions_nodes.SSL_GeminiInteractionsAPIKeyConfig.execute( api_key="secret123", api_version="v1beta", use_vertexai_env=False, vertexai_express=False, use_cache=True, cache_ttl_minutes=120, cache_seed=999, ) config = output[0] assert config["use_cache"] is True assert config["cache_ttl_minutes"] == 120 assert config["cache_seed"] == 999 def test_caching_with_cache_seed_and_usage(mock_interactions_client): mock_interactions_client.responses = [ {"text": "Cached response", "id": "int_cache_hit_1", "cached_tokens": 4096}, ] config = _config(use_cache=True, cache_seed=777) args = _execute_kwargs(config=config, use_seed=True, seed=123) out1 = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert out1[0] == "Cached response" assert out1[2] == 777 # actual_seed matches cache_seed when use_cache=True assert len(mock_interactions_client.calls) == 1 assert mock_interactions_client.calls[0]["generation_config"]["seed"] == 777 # Second call with identical params hits local result cache out2 = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert out2[0] == "Cached response" assert out2[2] == 777 assert len(mock_interactions_client.calls) == 1 def test_safety_config_node(): schema = interactions_nodes.SSL_GeminiInteractionsSafetyConfig.define_schema() assert schema.node_id == "SSL_GeminiInteractionsSafetyConfig" assert schema.category == "API/Gemini/Interactions" output = interactions_nodes.SSL_GeminiInteractionsSafetyConfig.execute( harassment="off", hate_speech="block_none", jailbreak="block_only_high", method="severity", ) settings = output[0] assert isinstance(settings, list) settings_dict = {s["type"]: s for s in settings} assert settings_dict["harassment"]["threshold"] == "off" assert settings_dict["harassment"]["method"] == "severity" assert settings_dict["hate_speech"]["threshold"] == "block_none" assert settings_dict["jailbreak"]["threshold"] == "block_only_high" def test_safety_preferences_off_and_method(mock_interactions_client): mock_interactions_client.responses = [{"text": "Unrestricted response", "id": "int_safe_1"}] args = _execute_kwargs(safety_level="off", safety_method="severity") output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert output[0] == "Unrestricted response" call = mock_interactions_client.calls[0] assert "safety_settings" in call assert len(call["safety_settings"]) == 10 assert all(s["threshold"] == "off" for s in call["safety_settings"]) assert all(s["method"] == "severity" for s in call["safety_settings"]) def test_safety_config_input_override(mock_interactions_client): mock_interactions_client.responses = [{"text": "Granular safety response", "id": "int_safe_2"}] custom_safety = [ {"type": "harassment", "threshold": "block_only_high"}, {"type": "jailbreak", "threshold": "off"}, ] args = _execute_kwargs( safety_level="block_none", # Should be overridden by safety_config safety_config=custom_safety, ) output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert output[0] == "Granular safety response" call = mock_interactions_client.calls[0] assert call["safety_settings"] == custom_safety def test_safety_block_error_reporting(mock_interactions_client): mock_interactions_client.responses = [{ "text": "", "id": "int_blocked", "errors": [SimpleNamespace(code="SAFETY", message="Content was blocked by safety policy")], "status": "failed", }] args = _execute_kwargs() output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args) assert output[0].startswith("API call/processing error:") assert "Content was blocked by safety policy" in output[0]