| 1 | import os |
| 2 | import sys |
| 3 | import unittest |
| 4 | from pathlib import Path |
| 5 | from unittest.mock import MagicMock, patch |
| 6 | |
| 7 | sys.path.insert(0, str(Path(__file__).parent.parent.parent)) |
| 8 | |
| 9 | from app.config import config |
| 10 | from app.services import twelvelabs |
| 11 | |
| 12 | RUN_INTEGRATION_TESTS = os.environ.get("MPT_RUN_INTEGRATION_TESTS", "").lower() in { |
| 13 | "1", |
| 14 | "true", |
| 15 | "yes", |
| 16 | } |
| 17 | |
| 18 | |
| 19 | class TestTwelveLabsService(unittest.TestCase): |
| 20 | """ |
| 21 | TwelveLabs 集成是完全 opt-in 的：未配置 twelvelabs_api_keys 时所有函数 |
| 22 | 都必须是无副作用的 no-op，行为与不接入 TwelveLabs 完全一致。 |
| 23 | 这些用例全部用 mock 替换 SDK 客户端，CI 不依赖真实网络或真实 API key。 |
| 24 | """ |
| 25 | |
| 26 | def setUp(self): |
| 27 | self.original_app_config = dict(config.app) |
| 28 | twelvelabs._embed_text_cached.cache_clear() |
| 29 | |
| 30 | def tearDown(self): |
| 31 | config.app.clear() |
| 32 | config.app.update(self.original_app_config) |
| 33 | twelvelabs._embed_text_cached.cache_clear() |
| 34 | |
| 35 | # ---------------- disabled / no-op behavior ---------------- |
| 36 | |
| 37 | def test_disabled_when_no_api_key(self): |
| 38 | config.app.pop("twelvelabs_api_keys", None) |
| 39 | self.assertFalse(twelvelabs.is_enabled()) |
| 40 | # rerank must return the input list unchanged |
| 41 | terms = ["b", "a", "c"] |
| 42 | self.assertEqual( |
| 43 | twelvelabs.rerank_terms_by_subject("subject", terms), terms |
| 44 | ) |
| 45 | # analyze must be a no-op returning None |
| 46 | self.assertIsNone(twelvelabs.analyze_clip("https://x/y.mp4")) |
| 47 | |
| 48 | def test_rerank_skipped_when_flag_off(self): |
| 49 | config.app["twelvelabs_api_keys"] = ["tlk_test"] |
| 50 | config.app["twelvelabs_rerank_terms"] = False |
| 51 | terms = ["b", "a"] |
| 52 | # Even enabled, with the flag off we must not touch order or call the API. |
| 53 | with patch.object(twelvelabs, "_client") as client: |
| 54 | result = twelvelabs.rerank_terms_by_subject("subject", terms) |
| 55 | self.assertEqual(result, terms) |
| 56 | client.assert_not_called() |
| 57 | |
| 58 | # ---------------- enabled rerank behavior ---------------- |
| 59 | |
| 60 | def _client_returning(self, vectors_by_text): |
| 61 | """Build a fake TwelveLabs client whose embed.create returns canned vectors.""" |
| 62 | |
| 63 | def fake_create(*, model_name, text): |
| 64 | seg = MagicMock() |
| 65 | seg.float_ = vectors_by_text[text] |
| 66 | resp = MagicMock() |
| 67 | resp.text_embedding.segments = [seg] |
| 68 | return resp |
| 69 | |
| 70 | client = MagicMock() |
| 71 | client.embed.create.side_effect = fake_create |
| 72 | return client |
| 73 | |
| 74 | def test_rerank_orders_by_cosine_to_subject(self): |
| 75 | config.app["twelvelabs_api_keys"] = ["tlk_test"] |
| 76 | config.app["twelvelabs_rerank_terms"] = True |
| 77 | |
| 78 | # subject aligned with "city"; "kitten" is orthogonal. |
| 79 | vectors = { |
| 80 | "city skyline": [1.0, 0.0, 0.0], |
| 81 | "downtown buildings": [0.9, 0.1, 0.0], # close to subject |
| 82 | "cute kitten": [0.0, 1.0, 0.0], # far from subject |
| 83 | } |
| 84 | client = self._client_returning(vectors) |
| 85 | |
| 86 | with patch.object(twelvelabs, "_client", return_value=client): |
| 87 | result = twelvelabs.rerank_terms_by_subject( |
| 88 | "city skyline", ["cute kitten", "downtown buildings"] |
| 89 | ) |
| 90 | |
| 91 | # most relevant term must come first |
| 92 | self.assertEqual(result, ["downtown buildings", "cute kitten"]) |
| 93 | |
| 94 | def test_rerank_falls_back_on_embed_failure(self): |
| 95 | config.app["twelvelabs_api_keys"] = ["tlk_test"] |
| 96 | config.app["twelvelabs_rerank_terms"] = True |
| 97 | |
| 98 | client = MagicMock() |
| 99 | client.embed.create.side_effect = RuntimeError("api down") |
| 100 | |
| 101 | terms = ["alpha", "beta"] |
| 102 | with patch.object(twelvelabs, "_client", return_value=client): |
| 103 | result = twelvelabs.rerank_terms_by_subject("subject", terms) |
| 104 | |
| 105 | # any failure must preserve the original order (never make things worse) |
| 106 | self.assertEqual(result, terms) |
| 107 | |
| 108 | def test_rerank_noop_for_single_term(self): |
| 109 | config.app["twelvelabs_api_keys"] = ["tlk_test"] |
| 110 | config.app["twelvelabs_rerank_terms"] = True |
| 111 | with patch.object(twelvelabs, "_client") as client: |
| 112 | result = twelvelabs.rerank_terms_by_subject("subject", ["only"]) |
| 113 | self.assertEqual(result, ["only"]) |
| 114 | client.assert_not_called() |
| 115 | |
| 116 | # ---------------- analyze_clip ---------------- |
| 117 | |
| 118 | def test_analyze_clip_returns_model_text(self): |
| 119 | config.app["twelvelabs_api_keys"] = ["tlk_test"] |
| 120 | |
| 121 | # analyze_clip() lazily imports `twelvelabs.types.VideoContext_Url`. |
| 122 | # The SDK is an optional extra, so the deterministic unit test must pass |
| 123 | # even without `uv sync --extra twelvelabs`. Inject lightweight stub |
| 124 | # modules so the internal import resolves; the mocked _client below does |
| 125 | # the rest. (When the real SDK *is* installed, these stubs are ignored.) |
| 126 | stub_types = type(sys)("twelvelabs.types") |
| 127 | stub_types.VideoContext_Url = lambda *, url: {"url": url} |
| 128 | stub_pkg = sys.modules.get("twelvelabs") or type(sys)("twelvelabs") |
| 129 | with patch.dict( |
| 130 | sys.modules, {"twelvelabs": stub_pkg, "twelvelabs.types": stub_types} |
| 131 | ): |
| 132 | self._run_analyze_clip_assertions() |
| 133 | |
| 134 | def _run_analyze_clip_assertions(self): |
| 135 | resp = MagicMock() |
| 136 | resp.data = "A city skyline at dusk." |
| 137 | client = MagicMock() |
| 138 | client.analyze.return_value = resp |
| 139 | |
| 140 | with patch.object(twelvelabs, "_client", return_value=client): |
| 141 | out = twelvelabs.analyze_clip( |
| 142 | "https://example.com/clip.mp4", prompt="describe" |
| 143 | ) |
| 144 | |
| 145 | self.assertEqual(out, "A city skyline at dusk.") |
| 146 | # max_tokens must be clamped to the Pegasus minimum (>=512) |
| 147 | self.assertGreaterEqual(client.analyze.call_args.kwargs["max_tokens"], 512) |
| 148 | |
| 149 | |
| 150 | @unittest.skipUnless( |
| 151 | RUN_INTEGRATION_TESTS and os.getenv("TWELVELABS_API_KEY"), |
| 152 | "live test: set MPT_RUN_INTEGRATION_TESTS=1 and TWELVELABS_API_KEY to run " |
| 153 | "against the real TwelveLabs API", |
| 154 | ) |
| 155 | class TestTwelveLabsLive(unittest.TestCase): |
| 156 | """Live contract check — only runs with MPT_RUN_INTEGRATION_TESTS=1 + a key.""" |
| 157 | |
| 158 | def setUp(self): |
| 159 | self.original_app_config = dict(config.app) |
| 160 | config.app["twelvelabs_api_keys"] = [os.environ["TWELVELABS_API_KEY"]] |
| 161 | config.app["twelvelabs_rerank_terms"] = True |
| 162 | twelvelabs._embed_text_cached.cache_clear() |
| 163 | |
| 164 | def tearDown(self): |
| 165 | config.app.clear() |
| 166 | config.app.update(self.original_app_config) |
| 167 | twelvelabs._embed_text_cached.cache_clear() |
| 168 | |
| 169 | def test_marengo_embedding_is_512_dim(self): |
| 170 | vec = twelvelabs.embed_text("a city skyline at night") |
| 171 | self.assertIsNotNone(vec) |
| 172 | self.assertEqual(len(vec), 512) |
| 173 | |
| 174 | def test_rerank_puts_relevant_term_first(self): |
| 175 | result = twelvelabs.rerank_terms_by_subject( |
| 176 | "city skyline at night", |
| 177 | ["cute kitten playing with yarn", "downtown buildings and traffic at dusk"], |
| 178 | ) |
| 179 | self.assertEqual(result[0], "downtown buildings and traffic at dusk") |
| 180 | |
| 181 | |
| 182 | if __name__ == "__main__": |
| 183 | unittest.main() |
| 184 |