import asyncio import base64 import types from io import BytesIO from unittest.mock import AsyncMock, patch import pytest from src import nodes class FakeVideoObject: def __init__(self, status, id="vid_1", error=None): self.status = status self.id = id self.error = error def _gemini_source(model="veo-3.1-fast-generate-preview", **overrides): settings = { "model": model, "size": "16:9", "resolution": "720", "seconds": 8, "person_generation": "allow_all", } settings.update(overrides) return {"source": "text", "provider": {"provider": "gemini", "model": settings}} def _openai_source(model="sora-2", **overrides): settings = {"model": model, "size": "16:9", "resolution": "1K", "seconds": 4} settings.update(overrides) return {"source": "text", "provider": {"provider": "openai", "model": settings}} @pytest.fixture(autouse=True) def _hidden(): nodes.APIVideoGenerate.hidden = types.SimpleNamespace(unique_id="test-node") yield nodes.APIVideoGenerate.hidden = None def test_generate_success_polls_until_completed(): async def run(): gen_mock = AsyncMock(return_value=FakeVideoObject("processing")) status_mock = AsyncMock( side_effect=[FakeVideoObject("processing"), FakeVideoObject("completed")] ) content_mock = AsyncMock(return_value=b"fake-mp4-bytes") with patch.object(nodes, "avideo_generation", gen_mock), patch.object( nodes, "avideo_status", status_mock ), patch.object(nodes, "avideo_content", content_mock), patch( "asyncio.sleep", AsyncMock(return_value=None) ): result = await nodes.APIVideoGenerate.execute( prompt="a cat riding a bike", source=_gemini_source(), ) assert isinstance(result, nodes.io.NodeOutput) video = result.args[0] assert isinstance(video, nodes.InputImpl.VideoFromFile) assert video.file_obj.getvalue() == b"fake-mp4-bytes" assert status_mock.call_count == 2 gen_mock.assert_awaited_once() assert gen_mock.call_args.kwargs["model"] == "gemini/veo-3.1-fast-generate-preview" assert gen_mock.call_args.kwargs["aspectRatio"] == "16:9" assert gen_mock.call_args.kwargs["resolution"] == "720p" assert gen_mock.call_args.kwargs["durationSeconds"] == 8 assert gen_mock.call_args.kwargs["personGeneration"] == "allow_all" assert "api_key" not in gen_mock.call_args.kwargs assert "api_key" not in status_mock.call_args.kwargs assert "api_key" not in content_mock.call_args.kwargs asyncio.run(run()) def test_api_key_threaded_through_every_call(): async def run(): gen_mock = AsyncMock(return_value=FakeVideoObject("completed")) status_mock = AsyncMock() content_mock = AsyncMock(return_value=b"bytes") with patch.object(nodes, "avideo_generation", gen_mock), patch.object( nodes, "avideo_status", status_mock ), patch.object(nodes, "avideo_content", content_mock), patch( "asyncio.sleep", AsyncMock(return_value=None) ): await nodes.APIVideoGenerate.execute( prompt="a dog surfing", source=_openai_source(), api_key="sk-test-123", ) assert gen_mock.call_args.kwargs["model"] == "openai/sora-2" assert gen_mock.call_args.kwargs["size"] == "1280x720" assert gen_mock.call_args.kwargs["seconds"] == "4" assert gen_mock.call_args.kwargs["api_key"] == "sk-test-123" assert content_mock.call_args.kwargs["api_key"] == "sk-test-123" status_mock.assert_not_awaited() # already completed, no polling needed asyncio.run(run()) def test_failed_generation_raises_runtime_error(): async def run(): gen_mock = AsyncMock(return_value=FakeVideoObject("processing")) status_mock = AsyncMock( return_value=FakeVideoObject("failed", error={"message": "content filtered"}) ) content_mock = AsyncMock() with patch.object(nodes, "avideo_generation", gen_mock), patch.object( nodes, "avideo_status", status_mock ), patch.object(nodes, "avideo_content", content_mock), patch( "asyncio.sleep", AsyncMock(return_value=None) ): with pytest.raises(RuntimeError, match="content filtered"): await nodes.APIVideoGenerate.execute( prompt="anything", source=_gemini_source(), ) content_mock.assert_not_awaited() asyncio.run(run()) def test_timeout_raises_when_generation_never_completes(): async def run(): gen_mock = AsyncMock(return_value=FakeVideoObject("processing")) status_mock = AsyncMock(return_value=FakeVideoObject("processing")) with patch.object(nodes, "avideo_generation", gen_mock), patch.object( nodes, "avideo_status", status_mock ), patch("asyncio.sleep", AsyncMock(return_value=None)): with pytest.raises(TimeoutError): await nodes.APIVideoGenerate.execute( prompt="anything", source=_gemini_source(), poll_interval=5.0, timeout=12.0, ) asyncio.run(run()) def test_runwayml_kwargs_mapping(): async def run(): gen_mock = AsyncMock(return_value=FakeVideoObject("completed")) content_mock = AsyncMock(return_value=b"bytes") source = { "source": "text", "provider": { "provider": "runwayml", "model": { "model": "gen3a_turbo", "size": "9:16", "resolution": "2K", "seconds": 10, "seed": 42, }, }, } with patch.object(nodes, "avideo_generation", gen_mock), patch.object( nodes, "avideo_content", content_mock ), patch("asyncio.sleep", AsyncMock(return_value=None)): await nodes.APIVideoGenerate.execute(prompt="anything", source=source) assert gen_mock.call_args.kwargs["model"] == "runwayml/gen3a_turbo" assert gen_mock.call_args.kwargs["ratio"] == "1080:1920" assert gen_mock.call_args.kwargs["duration"] == 10 assert gen_mock.call_args.kwargs["seed"] == 42 asyncio.run(run()) def test_schema_builds_nested_source_provider_model_combo(): schema = nodes.APIVideoGenerate.define_schema() source_input = next(i for i in schema.inputs if i.name == "source") assert [o.key for o in source_input.options] == ["text", "image"] text_option, image_option = source_input.options provider_input = text_option.inputs[0] assert provider_input.name == "provider" provider_keys = [o.key for o in provider_input.options] assert provider_keys == ["gemini", "openai", "azure", "runwayml"] for option in provider_input.options: model_input = option.inputs[0] assert model_input.name == "model" assert len(model_input.options) > 0 for model_option in model_input.options: field_names = {i.name for i in model_option.inputs} assert {"size", "resolution", "seconds"} <= field_names gemini_text_option = next(o for o in provider_input.options if o.key == "gemini") gemini_text_model_input = gemini_text_option.inputs[0] gemini_text_model_ids = {o.key for o in gemini_text_model_input.options} assert "veo-3.1-fast-generate-preview" in gemini_text_model_ids assert "gemini-omni-flash-preview" in gemini_text_model_ids veo_option = next(o for o in gemini_text_model_input.options if o.key == "veo-3.1-fast-generate-preview") veo_field_names = {i.name for i in veo_option.inputs} assert "person_generation" in veo_field_names omni_option = next(o for o in gemini_text_model_input.options if o.key == "gemini-omni-flash-preview") omni_field_names = {i.name for i in omni_option.inputs} assert "person_generation" not in omni_field_names image_provider_input = image_option.inputs[0] assert image_provider_input.name == "provider" assert [o.key for o in image_provider_input.options] == ["gemini"] gemini_model_input = image_provider_input.options[0].inputs[0] assert gemini_model_input.name == "model" gemini_image_model_ids = {o.key for o in gemini_model_input.options} assert "veo-3.1-fast-generate-preview" in gemini_image_model_ids assert "gemini-omni-flash-preview" in gemini_image_model_ids veo_image_option = next( o for o in gemini_model_input.options if o.key == "veo-3.1-fast-generate-preview" ) veo_image_field_names = {i.name for i in veo_image_option.inputs} assert {"size", "resolution", "seconds", "reference_image", "person_generation"} <= veo_image_field_names person_generation_input = next( i for i in veo_image_option.inputs if i.name == "person_generation" ) assert person_generation_input.kwargs["options"] == ["allow_adult"] omni_image_option = next( o for o in gemini_model_input.options if o.key == "gemini-omni-flash-preview" ) omni_image_field_names = {i.name for i in omni_image_option.inputs} assert {"size", "resolution", "seconds", "reference_image"} <= omni_image_field_names assert "person_generation" not in omni_image_field_names def test_gemini_image_to_video_kwargs_mapping(): async def run(): gen_mock = AsyncMock(return_value=FakeVideoObject("completed")) content_mock = AsyncMock(return_value=b"bytes") fake_tensor = object() fake_bytesio = object() source = { "source": "image", "provider": { "provider": "gemini", "model": { "model": "veo-3.1-fast-generate-preview", "size": "9:16", "resolution": "1080", "seconds": 8, "reference_image": fake_tensor, "person_generation": "allow_adult", }, }, } with patch.object(nodes, "avideo_generation", gen_mock), patch.object( nodes, "avideo_content", content_mock ), patch.object( nodes.GeminiModel, "image_tensor_to_bytesio", return_value=fake_bytesio ) as tensor_mock, patch("asyncio.sleep", AsyncMock(return_value=None)): await nodes.APIVideoGenerate.execute(prompt="animate this photo", source=source) tensor_mock.assert_called_once_with(fake_tensor) kwargs = gen_mock.call_args.kwargs assert kwargs["model"] == "gemini/veo-3.1-fast-generate-preview" assert kwargs["aspectRatio"] == "9:16" assert kwargs["resolution"] == "1080p" assert kwargs["durationSeconds"] == 8 assert kwargs["personGeneration"] == "allow_adult" assert kwargs["input_reference"] is fake_bytesio asyncio.run(run()) def test_gemini_image_to_video_requires_reference_image(): async def run(): gen_mock = AsyncMock(return_value=FakeVideoObject("completed")) source = { "source": "image", "provider": { "provider": "gemini", "model": { "model": "veo-3.1-fast-generate-preview", "size": "9:16", "resolution": "1080", "seconds": 8, "person_generation": "allow_adult", }, }, } with patch.object(nodes, "avideo_generation", gen_mock): with pytest.raises(ValueError, match="reference_image"): await nodes.APIVideoGenerate.execute(prompt="animate this photo", source=source) gen_mock.assert_not_awaited() asyncio.run(run()) class FakeInteractionResponse: def __init__(self, steps): self.steps = steps def _gemini_omni_source(model="gemini-omni-flash-preview", **overrides): settings = {"model": model, "size": "16:9", "resolution": "720", "seconds": 5} settings.update(overrides) return {"source": "text", "provider": {"provider": "gemini", "model": settings}} def test_gemini_omni_uses_interactions_api_not_polling(): async def run(): interaction_mock = AsyncMock( return_value=FakeInteractionResponse( steps=[ {"type": "user_input", "content": [{"type": "text", "text": "hi"}]}, { "type": "model_output", "content": [ { "type": "video", "data": base64.b64encode(b"fake-mp4-bytes").decode(), "mime_type": "video/mp4", } ], }, ] ) ) gen_mock = AsyncMock() with patch.object(nodes, "acreate_interaction", interaction_mock), patch.object( nodes, "avideo_generation", gen_mock ): result = await nodes.APIVideoGenerate.execute( prompt="a marble rolling down a track", source=_gemini_omni_source(), ) gen_mock.assert_not_awaited() interaction_mock.assert_awaited_once() kwargs = interaction_mock.call_args.kwargs assert kwargs["model"] == "gemini/gemini-omni-flash-preview" assert kwargs["input"] == "a marble rolling down a track" assert kwargs["response_format"] == { "type": "video", "aspect_ratio": "16:9", "resolution": "720p", } assert "api_key" not in kwargs video = result.args[0] assert isinstance(video, nodes.InputImpl.VideoFromFile) assert video.file_obj.getvalue() == b"fake-mp4-bytes" asyncio.run(run()) def test_gemini_omni_uri_delivery_raises_clear_error(): async def run(): interaction_mock = AsyncMock( return_value=FakeInteractionResponse( steps=[ { "type": "model_output", "content": [{"type": "video", "uri": "https://example.com/video.mp4"}], } ] ) ) with patch.object(nodes, "acreate_interaction", interaction_mock): with pytest.raises(RuntimeError, match="uri"): await nodes.APIVideoGenerate.execute( prompt="anything", source=_gemini_omni_source(), ) asyncio.run(run()) def test_gemini_omni_image_to_video_sends_text_and_image_content(): async def run(): interaction_mock = AsyncMock( return_value=FakeInteractionResponse( steps=[ { "type": "model_output", "content": [ { "type": "video", "data": base64.b64encode(b"fake-mp4-bytes").decode(), } ], } ] ) ) fake_tensor = object() fake_bytesio_content = b"fake-png-bytes" source = { "source": "image", "provider": { "provider": "gemini", "model": { "model": "gemini-omni-flash-preview", "size": "16:9", "resolution": "720", "seconds": 5, "reference_image": fake_tensor, }, }, } with patch.object(nodes, "acreate_interaction", interaction_mock), patch.object( nodes.GeminiModel, "image_tensor_to_bytesio_list", return_value=[BytesIO(fake_bytesio_content)], ) as tensor_mock: result = await nodes.APIVideoGenerate.execute( prompt="make the person in the photo wave", source=source, ) tensor_mock.assert_called_once_with(fake_tensor) kwargs = interaction_mock.call_args.kwargs assert kwargs["model"] == "gemini/gemini-omni-flash-preview" assert kwargs["input"] == [ {"type": "text", "text": "make the person in the photo wave"}, { "type": "image", "data": base64.b64encode(fake_bytesio_content).decode(), "mime_type": "image/png", }, ] assert kwargs["response_format"] == { "type": "video", "aspect_ratio": "16:9", "resolution": "720p", } video = result.args[0] assert video.file_obj.getvalue() == b"fake-mp4-bytes" asyncio.run(run()) def test_gemini_omni_image_to_video_sends_one_content_item_per_batched_image(): async def run(): interaction_mock = AsyncMock( return_value=FakeInteractionResponse( steps=[ { "type": "model_output", "content": [{"type": "video", "data": base64.b64encode(b"vid").decode()}], } ] ) ) fake_tensor = object() frame_bytes = [b"frame-1", b"frame-2", b"frame-3"] source = { "source": "image", "provider": { "provider": "gemini", "model": { "model": "gemini-omni-flash-preview", "size": "16:9", "resolution": "720", "seconds": 5, "reference_image": fake_tensor, }, }, } with patch.object(nodes, "acreate_interaction", interaction_mock), patch.object( nodes.GeminiModel, "image_tensor_to_bytesio_list", return_value=[BytesIO(b) for b in frame_bytes], ): await nodes.APIVideoGenerate.execute(prompt="combine these", source=source) kwargs = interaction_mock.call_args.kwargs image_items = [item for item in kwargs["input"] if item["type"] == "image"] assert len(image_items) == 3 assert [item["data"] for item in image_items] == [ base64.b64encode(b).decode() for b in frame_bytes ] asyncio.run(run()) def test_veo_image_to_video_rejects_batched_reference_image(): async def run(): gen_mock = AsyncMock() fake_tensor = types.SimpleNamespace(shape=(3, 64, 64, 3)) source = { "source": "image", "provider": { "provider": "gemini", "model": { "model": "veo-3.1-fast-generate-preview", "size": "16:9", "resolution": "720", "seconds": 8, "reference_image": fake_tensor, "person_generation": "allow_adult", }, }, } with patch.object(nodes, "avideo_generation", gen_mock): with pytest.raises(ValueError, match="single reference image"): await nodes.APIVideoGenerate.execute(prompt="anything", source=source) gen_mock.assert_not_awaited() asyncio.run(run()) def test_azure_provider_string_uses_azure_prefix(): async def run(): gen_mock = AsyncMock(return_value=FakeVideoObject("completed")) content_mock = AsyncMock(return_value=b"bytes") source = _openai_source(model="sora-2-pro") source["provider"]["provider"] = "azure" with patch.object(nodes, "avideo_generation", gen_mock), patch.object( nodes, "avideo_content", content_mock ), patch("asyncio.sleep", AsyncMock(return_value=None)): await nodes.APIVideoGenerate.execute(prompt="anything", source=source) assert gen_mock.call_args.kwargs["model"] == "azure/sora-2-pro" asyncio.run(run())