mirror of
https://github.com/TheSmallHanCat/flow2api.git
synced 2026-05-06 22:13:48 +08:00
214 lines
8.2 KiB
Python
214 lines
8.2 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")
|
|
|
|
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()
|