mirror of
https://github.com/TheSmallHanCat/flow2api.git
synced 2026-05-22 20:31:42 +08:00
316 lines
12 KiB
Python
316 lines
12 KiB
Python
import types
|
|
import unittest
|
|
from unittest.mock import AsyncMock
|
|
|
|
from src.core.model_resolver import resolve_model_name
|
|
from src.services.flow_client import FlowClient
|
|
from src.services.generation_handler import MODEL_CONFIG, GenerationHandler
|
|
|
|
|
|
class VeoLiteModelResolverTests(unittest.TestCase):
|
|
def test_resolve_t2v_lite_alias_to_portrait_variant(self):
|
|
request = types.SimpleNamespace(
|
|
generationConfig=types.SimpleNamespace(aspectRatio="portrait")
|
|
)
|
|
|
|
resolved = resolve_model_name(
|
|
"veo_3_1_t2v_lite",
|
|
request=request,
|
|
model_config=MODEL_CONFIG,
|
|
)
|
|
|
|
self.assertEqual(resolved, "veo_3_1_t2v_lite_portrait")
|
|
|
|
def test_resolve_quality_4s_upsample_alias_to_portrait_variant(self):
|
|
request = types.SimpleNamespace(
|
|
generationConfig=types.SimpleNamespace(aspectRatio="portrait")
|
|
)
|
|
|
|
resolved = resolve_model_name(
|
|
"veo_3_1_t2v_4s_4k",
|
|
request=request,
|
|
model_config=MODEL_CONFIG,
|
|
)
|
|
|
|
self.assertEqual(resolved, "veo_3_1_t2v_portrait_4s_4k")
|
|
|
|
def test_resolve_video_image_size_to_upsample_variant(self):
|
|
request = types.SimpleNamespace(
|
|
generationConfig=types.SimpleNamespace(aspectRatio="landscape", imageSize="1080p")
|
|
)
|
|
|
|
resolved = resolve_model_name(
|
|
"veo_3_1_i2v_s_6s",
|
|
request=request,
|
|
model_config=MODEL_CONFIG,
|
|
)
|
|
|
|
self.assertEqual(resolved, "veo_3_1_i2v_s_6s_1080p")
|
|
|
|
|
|
class VeoLiteGenerationHandlerTests(unittest.TestCase):
|
|
def test_tier_two_does_not_upgrade_lite_model_to_fake_ultra(self):
|
|
handler = GenerationHandler.__new__(GenerationHandler)
|
|
|
|
model_key, message = handler._resolve_video_model_key_for_tier(
|
|
{
|
|
"model_key": "veo_3_1_t2v_lite",
|
|
"allow_tier_upgrade": False,
|
|
},
|
|
"PAYGATE_TIER_TWO",
|
|
)
|
|
|
|
self.assertEqual(model_key, "veo_3_1_t2v_lite")
|
|
self.assertIsNone(message)
|
|
|
|
def test_tier_two_still_upgrades_regular_model(self):
|
|
handler = GenerationHandler.__new__(GenerationHandler)
|
|
|
|
model_key, message = handler._resolve_video_model_key_for_tier(
|
|
{
|
|
"model_key": "veo_3_1_t2v_fast",
|
|
},
|
|
"PAYGATE_TIER_TWO",
|
|
)
|
|
|
|
self.assertEqual(model_key, "veo_3_1_t2v_fast_ultra")
|
|
self.assertIn("ultra", message)
|
|
|
|
def test_quality_model_does_not_upgrade_to_fake_ultra(self):
|
|
handler = GenerationHandler.__new__(GenerationHandler)
|
|
|
|
model_key, message = handler._resolve_video_model_key_for_tier(
|
|
{
|
|
"model_key": "veo_3_1_t2v",
|
|
},
|
|
"PAYGATE_TIER_TWO",
|
|
)
|
|
|
|
self.assertEqual(model_key, "veo_3_1_t2v")
|
|
self.assertIsNone(message)
|
|
|
|
def test_quality_4s_upsample_model_generates_then_upsamples(self):
|
|
cfg = MODEL_CONFIG["veo_3_1_t2v_4s_4k"]
|
|
|
|
self.assertEqual(cfg["model_key"], "veo_3_1_t2v_quality_4s")
|
|
self.assertEqual(cfg["video_type"], "t2v")
|
|
self.assertEqual(cfg["upsample"]["model_key"], "veo_3_1_upsampler_4k")
|
|
self.assertEqual(cfg["upsample"]["resolution"], "VIDEO_RESOLUTION_4K")
|
|
|
|
def test_quality_6s_i2v_1080p_model_generates_then_upsamples(self):
|
|
cfg = MODEL_CONFIG["veo_3_1_i2v_s_6s_1080p"]
|
|
|
|
self.assertEqual(cfg["model_key"], "veo_3_1_i2v_s_quality_6s_fl")
|
|
self.assertEqual(cfg["video_type"], "i2v")
|
|
self.assertEqual(cfg["upsample"]["model_key"], "veo_3_1_upsampler_1080p")
|
|
self.assertEqual(cfg["upsample"]["resolution"], "VIDEO_RESOLUTION_1080P")
|
|
|
|
def test_short_duration_models_include_explicit_landscape_aliases(self):
|
|
expected_aliases = {
|
|
"veo_3_1_t2v_landscape_4s": "veo_3_1_t2v_4s",
|
|
"veo_3_1_t2v_landscape_6s": "veo_3_1_t2v_6s",
|
|
"veo_3_1_i2v_s_landscape_4s": "veo_3_1_i2v_s_4s",
|
|
"veo_3_1_i2v_s_landscape_6s": "veo_3_1_i2v_s_6s",
|
|
"veo_3_1_t2v_landscape_4s_4k": "veo_3_1_t2v_4s_4k",
|
|
"veo_3_1_i2v_s_landscape_6s_1080p": "veo_3_1_i2v_s_6s_1080p",
|
|
}
|
|
|
|
for alias, target in expected_aliases.items():
|
|
self.assertIn(alias, MODEL_CONFIG)
|
|
self.assertEqual(MODEL_CONFIG[alias], MODEL_CONFIG[target])
|
|
|
|
def test_r2v_models_include_explicit_landscape_aliases(self):
|
|
expected_aliases = {
|
|
"veo_3_1_r2v_fast_landscape": "veo_3_1_r2v_fast",
|
|
"veo_3_1_r2v_fast_landscape_ultra": "veo_3_1_r2v_fast_ultra",
|
|
"veo_3_1_r2v_fast_landscape_ultra_relaxed": "veo_3_1_r2v_fast_ultra_relaxed",
|
|
"veo_3_1_r2v_fast_landscape_ultra_4k": "veo_3_1_r2v_fast_ultra_4k",
|
|
"veo_3_1_r2v_fast_landscape_ultra_1080p": "veo_3_1_r2v_fast_ultra_1080p",
|
|
}
|
|
|
|
for alias, target in expected_aliases.items():
|
|
self.assertIn(alias, MODEL_CONFIG)
|
|
self.assertEqual(MODEL_CONFIG[alias], MODEL_CONFIG[target])
|
|
|
|
def test_direct_upsampler_keys_are_not_public_models(self):
|
|
self.assertNotIn("veo_3_1_upsampler_4k", MODEL_CONFIG)
|
|
self.assertNotIn("veo_3_1_upsampler_1080p", MODEL_CONFIG)
|
|
|
|
|
|
class VeoLiteFlowClientTests(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
self.client = FlowClient(proxy_manager=None)
|
|
self.client._acquire_video_launch_gate = AsyncMock(return_value=(True, None, None))
|
|
self.client._release_video_launch_gate = AsyncMock()
|
|
self.client._get_recaptcha_token = AsyncMock(return_value=("recaptcha-token", "browser-1"))
|
|
self.client._notify_browser_captcha_request_finished = AsyncMock()
|
|
|
|
async def test_generate_video_text_uses_v2_payload_for_lite(self):
|
|
captured = {}
|
|
|
|
async def fake_make_request(method, url, json_data, use_at, at_token, **kwargs):
|
|
captured["url"] = url
|
|
captured["json_data"] = json_data
|
|
return {"operations": [{"operation": {"name": "task-1"}}]}
|
|
|
|
self.client._make_request = AsyncMock(side_effect=fake_make_request)
|
|
|
|
await self.client.generate_video_text(
|
|
at="at-token",
|
|
project_id="project-1",
|
|
prompt="猫猫",
|
|
model_key="veo_3_1_t2v_lite",
|
|
aspect_ratio="VIDEO_ASPECT_RATIO_LANDSCAPE",
|
|
use_v2_model_config=True,
|
|
)
|
|
|
|
json_data = captured["json_data"]
|
|
request_data = json_data["requests"][0]
|
|
self.assertTrue(json_data["useV2ModelConfig"])
|
|
self.assertIn("batchId", json_data["mediaGenerationContext"])
|
|
self.assertEqual(
|
|
request_data["textInput"]["structuredPrompt"]["parts"][0]["text"],
|
|
"猫猫",
|
|
)
|
|
self.assertNotIn("prompt", request_data["textInput"])
|
|
self.assertEqual(request_data["videoModelKey"], "veo_3_1_t2v_lite")
|
|
self.assertEqual(
|
|
json_data["mediaGenerationContext"]["audioFailurePreference"],
|
|
"BLOCK_SILENCED_VIDEOS",
|
|
)
|
|
|
|
async def test_generate_video_text_normalizes_media_only_create_response(self):
|
|
captured = {}
|
|
|
|
async def fake_make_request(method, url, json_data, use_at, at_token, **kwargs):
|
|
captured["json_data"] = json_data
|
|
return {
|
|
"remainingCredits": 30,
|
|
"workflows": [
|
|
{
|
|
"name": "workflow-1",
|
|
"metadata": {"primaryMediaId": "media-1"},
|
|
"projectId": "project-1",
|
|
}
|
|
],
|
|
"media": [
|
|
{
|
|
"name": "media-1",
|
|
"projectId": "project-1",
|
|
"mediaMetadata": {
|
|
"mediaStatus": {
|
|
"mediaGenerationStatus": "MEDIA_GENERATION_STATUS_PENDING"
|
|
}
|
|
},
|
|
}
|
|
],
|
|
}
|
|
|
|
self.client._make_request = AsyncMock(side_effect=fake_make_request)
|
|
|
|
result = await self.client.generate_video_text(
|
|
at="at-token",
|
|
project_id="project-1",
|
|
prompt="猫猫",
|
|
model_key="veo_3_1_t2v_lite",
|
|
aspect_ratio="VIDEO_ASPECT_RATIO_LANDSCAPE",
|
|
use_v2_model_config=True,
|
|
)
|
|
|
|
self.assertEqual(
|
|
captured["json_data"]["mediaGenerationContext"]["audioFailurePreference"],
|
|
"BLOCK_SILENCED_VIDEOS",
|
|
)
|
|
self.assertEqual(result["operations"][0]["operation"]["name"], "media-1")
|
|
self.assertEqual(result["operations"][0]["projectId"], "project-1")
|
|
self.assertEqual(
|
|
result["operations"][0]["status"],
|
|
"MEDIA_GENERATION_STATUS_PENDING",
|
|
)
|
|
|
|
async def test_check_video_status_uses_media_payload_and_normalizes_response(self):
|
|
captured = {}
|
|
|
|
async def fake_make_request(method, url, json_data, use_at, at_token, **kwargs):
|
|
captured["json_data"] = json_data
|
|
return {
|
|
"media": [
|
|
{
|
|
"name": "media-1",
|
|
"projectId": "project-1",
|
|
"mediaMetadata": {
|
|
"mediaStatus": {
|
|
"mediaGenerationStatus": "MEDIA_GENERATION_STATUS_SUCCESSFUL"
|
|
}
|
|
},
|
|
"video": {
|
|
"fifeUrl": "https://flow-content.google/video/11111111-1111-1111-1111-111111111111?token=abc",
|
|
"generatedVideo": {
|
|
"aspectRatio": "VIDEO_ASPECT_RATIO_LANDSCAPE"
|
|
},
|
|
},
|
|
}
|
|
]
|
|
}
|
|
|
|
self.client._make_request = AsyncMock(side_effect=fake_make_request)
|
|
|
|
result = await self.client.check_video_status(
|
|
at="at-token",
|
|
operations=[
|
|
{
|
|
"operation": {"name": "media-1"},
|
|
"projectId": "project-1",
|
|
}
|
|
],
|
|
)
|
|
|
|
self.assertEqual(
|
|
captured["json_data"],
|
|
{"media": [{"name": "media-1", "projectId": "project-1"}]},
|
|
)
|
|
operation = result["operations"][0]
|
|
self.assertEqual(operation["operation"]["name"], "media-1")
|
|
self.assertEqual(operation["status"], "MEDIA_GENERATION_STATUS_SUCCESSFUL")
|
|
self.assertEqual(
|
|
operation["operation"]["metadata"]["video"]["fifeUrl"],
|
|
"https://flow-content.google/video/11111111-1111-1111-1111-111111111111?token=abc",
|
|
)
|
|
|
|
async def test_generate_video_start_end_uses_v2_payload_for_interpolation_lite(self):
|
|
captured = {}
|
|
|
|
async def fake_make_request(method, url, json_data, use_at, at_token, **kwargs):
|
|
captured["url"] = url
|
|
captured["json_data"] = json_data
|
|
return {"operations": [{"operation": {"name": "task-2"}}]}
|
|
|
|
self.client._make_request = AsyncMock(side_effect=fake_make_request)
|
|
|
|
await self.client.generate_video_start_end(
|
|
at="at-token",
|
|
project_id="project-1",
|
|
prompt="变身猫猫",
|
|
model_key="veo_3_1_interpolation_lite",
|
|
aspect_ratio="VIDEO_ASPECT_RATIO_PORTRAIT",
|
|
start_media_id="start-media",
|
|
end_media_id="end-media",
|
|
use_v2_model_config=True,
|
|
)
|
|
|
|
json_data = captured["json_data"]
|
|
request_data = json_data["requests"][0]
|
|
self.assertTrue(json_data["useV2ModelConfig"])
|
|
self.assertIn("batchId", json_data["mediaGenerationContext"])
|
|
self.assertEqual(request_data["videoModelKey"], "veo_3_1_interpolation_lite")
|
|
self.assertEqual(request_data["startImage"]["mediaId"], "start-media")
|
|
self.assertEqual(request_data["endImage"]["mediaId"], "end-media")
|
|
self.assertEqual(
|
|
request_data["textInput"]["structuredPrompt"]["parts"][0]["text"],
|
|
"变身猫猫",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|