import contextlib import copy import io import json import os from pathlib import Path import subprocess import sys import tempfile import unittest from unittest.mock import patch import wave import diarization def word(start, end, speaker=9, text="speech"): return {"word": text, "start": start, "end": end, "speaker": speaker} def transcript(*words): return {"text": " ".join(item["word"] for item in words), "words": list(words)} def turn(start, end, speaker="voice-A", **extra): return dict(start=start, end=end, speaker=speaker, **extra) def job(*segments, **extra): return dict(jobId="test-job", status="succeeded", output={"exclusiveDiarization": list(segments)}, **extra) def align(source, result, speakers=None): with contextlib.redirect_stderr(io.StringIO()): return diarization.align(source, result, speakers) def write_audio(directory, duration=30): path = Path(directory) / "audio.wav" with wave.open(str(path), "wb") as audio: audio.setparams((1, 2, 16000, 0, "NONE", "not compressed")) audio.writeframes(b"\0\0" * int(duration * 16000)) return path class AlignmentTests(unittest.TestCase): def test_global_timeline_reunites_host_in_ads_and_across_request_local_ids(self): source = transcript(word(1, 3, 7, "Host interview."), word(4, 6, 8, "Guest answer."), word(601, 603, 8, "Host ad with a different microphone."), word(1001, 1003, 9, "Host interview again."), word(1004, 1006, 7, "Guest answer again.")) timeline = job(turn(0, 3.5, "HOST"), turn(4, 8, "GUEST"), turn(600, 650, "HOST"), turn(1000, 1003.5, "HOST"), turn(1004, 1020, "GUEST")) result = align(source, timeline, speakers=2) self.assertEqual([w["speaker"] for w in result["words"]], [0, 1, 0, 0, 1]) def test_reused_provider_local_id_never_merges_distinct_global_voices(self): source = transcript(word(1, 2, 0), word(4, 5, 0)) result = align(source, job(turn(0, 3, "a"), turn(4, 6, "b"))) self.assertEqual([w["speaker"] for w in result["words"]], [0, 1]) def test_all_original_text_order_timestamps_and_other_fields_survive(self): source = transcript(word(8, 12, None, " Original\t text "), word(7, 9, 3, "Earlier timing")) source["words"][0]["confidence"] = 0.7 source["words"][1].pop("speaker") source["metadata"] = {"nested": [1, 2]} before = copy.deepcopy(source) result = align(source, job(turn(0, 20))) self.assertEqual(source, before) self.assertEqual({k: v for k, v in result.items() if k != "words"}, {k: v for k, v in before.items() if k != "words"}) for original, changed in zip(before["words"], result["words"]): self.assertEqual({k: v for k, v in original.items() if k != "speaker"}, {k: v for k, v in changed.items() if k != "speaker"}) self.assertEqual([w["speaker"] for w in result["words"]], [0, 0]) def test_ids_follow_first_assigned_nonempty_word_not_untranscribed_timeline_turns(self): source = transcript(word(5, 6, text=" "), word(8, 9, text="First"), word(4, 4.5, text="Second")) result = align(source, job(turn(0, 2, "ignored"), turn(3, 6, "later"), turn(7, 10, "first"))) self.assertEqual([w["speaker"] for w in result["words"]], [None, 0, 1]) def test_overlap_aggregates_same_speaker_across_separate_turns(self): result = align(transcript(word(0, 10)), job(turn(0, 4, "a"), turn(4, 6, "b"), turn(6, 10, "a"))) self.assertEqual(result["words"][0]["speaker"], 0) def test_dominant_overlap_wins_but_equal_voices_or_low_coverage_stay_unknown(self): for segments, interval, expected in [ ([turn(0, 6, "a"), turn(6, 10, "b")], (0, 10), 0), ([turn(0, 5, "a"), turn(5, 10, "b")], (0, 10), None), ([turn(0, 1, "a")], (0, 10), None), ]: with self.subTest(segments=segments): self.assertEqual(align(transcript(word(*interval)), job(*segments))["words"][0]["speaker"], expected) def test_nearest_is_bounded_for_whole_word_and_not_only_closest_edge(self): source = transcript(word(1.01, 1.2), word(1.01, 2), word(0, 0.1), word(8, 8.1)) log = io.StringIO() with contextlib.redirect_stderr(log): result = diarization.align(source, job(turn(0.2, 1))) self.assertEqual([w["speaker"] for w in result["words"]], [0, None, 0, None]) self.assertIn("2 word(s) used a nearest", log.getvalue()) self.assertIn("2 word(s) have ambiguous", log.getvalue()) def test_zero_duration_and_nearest_ties_do_not_force_a_speaker(self): cases = [ (job(turn(0, 1, "a"), turn(1, 2, "b")), 1, None), (job(turn(0, 1, "a"), turn(1.2, 2, "b")), 1.1, None), (job(turn(0, 1, "a"), turn(1.2, 2, "a")), 1.1, 0), (job(turn(0, 1, "a")), 0.5, 0), (job(turn(0, 1, "a")), 1.3, None), ] for timeline, timestamp, expected in cases: with self.subTest(timestamp=timestamp, timeline=timeline): self.assertEqual(align(transcript(word(timestamp, timestamp)), timeline)["words"][0]["speaker"], expected) def test_empty_timeline_is_valid_only_without_transcribed_speech(self): self.assertEqual(align(transcript(), job()), transcript()) with self.assertRaisesRegex(ValueError, "empty speaker timeline"): align(transcript(word(0, 1, 42)), job()) def test_known_count_warns_without_merging_or_fabricating_speakers(self): log = io.StringIO() with contextlib.redirect_stderr(log): result = diarization.align(transcript(word(0, 1), word(2, 3), word(4, 5)), job(turn(0, 1, "a"), turn(2, 3, "b"), turn(4, 5, "c")), 2) self.assertEqual([w["speaker"] for w in result["words"]], [0, 1, 2]) self.assertIn("expected 2 speakers", log.getvalue()) def test_malformed_incomplete_unsorted_or_overlapping_timelines_fail(self): invalid = [None, {}, {"status": "running"}, dict(job(), status="failed"), dict(job(), output={}), dict(job(), error="failure"), dict(job(), output={"error": "failure", "exclusiveDiarization": []}), job(None), job(turn(-1, 1)), job(turn(2, 1)), job(turn(1, 1)), job(turn(True, 2)), job(turn(0, float("nan"))), job(turn(0, float("inf"))), job(turn(0, 1, None)), job(turn(0, 1, " ")), job(turn(0, 1, 0)), job(turn(0, 3), turn(2, 4)), job(turn(5, 6), turn(1, 2)), job(turn(0, 1, confidence=0.7)), job(turn(0, 1, confidence={"a": 101})), job(turn(0, 1, confidence={"a": True})), job(turn(0, 1, confidence={"a": float("nan")}))] for timeline in invalid: with self.subTest(timeline=timeline), self.assertRaises(ValueError): align(transcript(word(0, 1)), timeline) align(transcript(word(0, 1)), job(turn(0, 1, confidence={"voice-A": 97.2, "voice-B": 2.8}))) def test_invalid_words_or_known_speaker_counts_fail(self): for item in (word(-1, 1), word(1, 0), word(False, 1), word(0, float("nan"))): with self.subTest(item=item), self.assertRaises(ValueError): align(transcript(item), job(turn(0, 2))) for count in (0, -1, True, 1.5): with self.subTest(count=count), self.assertRaises(ValueError): align(transcript(word(0, 1)), job(turn(0, 2)), count) def test_cli_output_only_after_complete_validation(self): with tempfile.TemporaryDirectory() as directory: source = Path(directory) / "transcript.json" result = Path(directory) / "result.json" source.write_text(json.dumps(transcript(word(0, 1, text="Must survive.")))) result.write_text(json.dumps(job(turn(0, 2)))) command = [sys.executable, "-B", str(Path(diarization.__file__)), "apply", str(source), str(result)] completed = subprocess.run([*command, "--json"], capture_output=True, text=True) self.assertEqual(completed.returncode, 0, completed.stderr) self.assertEqual(json.loads(completed.stdout)["words"][0]["speaker"], 0) completed = subprocess.run([*command, "--srt"], capture_output=True, text=True) self.assertEqual(completed.returncode, 0, completed.stderr) self.assertIn("00:00:00,000 --> 00:00:01,000", completed.stdout) result.write_text('{"status":"succeeded","status":"running"}') completed = subprocess.run(command, capture_output=True, text=True) self.assertNotEqual(completed.returncode, 0) self.assertEqual(completed.stdout, "") self.assertIn("Duplicate JSON key", completed.stderr) class FakeTransport: """Exercise real Client command construction without executing curl.""" def __init__(self, polls=None, failure=None): self.polls = list(polls or [job(turn(0, 10, confidence={"voice-A": 99}))]) self.failure = failure self.calls = [] self.submissions = [] self.uploads = [] self.authenticated = [] def __call__(self, command, **kwargs): self.calls.append(command) if command[0] == "ffmpeg": Path(command[-1]).write_bytes(b"fLaC fake whole recording") return subprocess.CompletedProcess(command, 0, stderr=b"") output = Path(command[command.index("-o") + 1]) method = command[command.index("--request") + 1] url = command[-1] if method != "PUT": header = Path(command[command.index("--header") + 1][1:]) assert header.stat().st_mode & 0o777 == 0o600 self.authenticated.append(header.read_text()) if self.failure and self.failure[0] in url: status, body, code = self.failure[1:] elif url.endswith("/media/input"): self.media_request = json.loads(Path(command[command.index("--data-binary") + 1][1:]).read_text()) status, body, code = 200, {"url": "https://storage.example/file?secret=signature"}, 0 elif method == "PUT": self.uploads.append(Path(command[command.index("--upload-file") + 1]).read_bytes()) status, body, code = 200, "", 0 elif url.endswith("/diarize"): self.submissions.append(json.loads(Path(command[command.index("--data-binary") + 1][1:]).read_text())) status, body, code = 200, {"jobId": "test-job", "status": "created"}, 0 else: item = self.polls.pop(0) status, body, code = item if isinstance(item, tuple) else (200, item, 0) output.write_text(body if isinstance(body, str) else json.dumps(body)) return subprocess.CompletedProcess(command, code, stdout=str(status).encode(), stderr=b"sensitive https://storage.example/file?secret=signature") class RunTests(unittest.TestCase): def execute(self, fake, directory, **kwargs): audio = write_audio(directory) with patch.dict(os.environ, {"PYANNOTE_API_KEY": "test-secret-key"}), \ patch.object(diarization.subprocess, "run", side_effect=fake), \ patch.object(diarization.time, "sleep"), contextlib.redirect_stderr(io.StringIO()) as log: result = diarization.run(audio, Path(directory) / "work", **kwargs) return result, log.getvalue() def test_complete_flac_upload_exact_model_flags_and_private_request_cleanup(self): fake = FakeTransport() with tempfile.TemporaryDirectory() as directory: result, _ = self.execute(fake, directory, speakers=2) work = Path(directory) / "work" self.assertEqual(json.loads((work / "result.json").read_text()), result) self.assertEqual([p.name for p in work.iterdir()], ["result.json"]) self.assertEqual(fake.uploads, [b"fLaC fake whole recording"]) self.assertEqual(fake.submissions, [dict(url=fake.media_request["url"], model="precision-2", exclusive=True, turnLevelConfidence=True, transcription=False, numSpeakers=2)]) self.assertRegex(fake.media_request["url"], r"^media://[a-f0-9-]+\.flac$") encoding = fake.calls[0] self.assertNotIn("-ss", encoding) self.assertNotIn("-t", encoding) self.assertIn("flac", encoding) for command in fake.calls[1:]: self.assertIn("--max-time", command) self.assertIn("--connect-timeout", command) self.assertIn("--proto", command) self.assertNotIn("--location", command) self.assertNotIn("test-secret-key", " ".join(command)) self.assertEqual(any(value.startswith("@") for value in command if "authorization.header" in value), command[command.index("--request") + 1] != "PUT") self.assertEqual(fake.authenticated, ["Authorization: Bearer test-secret-key\n"] * 3) def test_hosted_community_override_is_sent_unchanged(self): fake = FakeTransport() with tempfile.TemporaryDirectory() as directory: self.execute(fake, directory, model="community-1") self.assertEqual(fake.submissions[0]["model"], "community-1") def test_saved_result_keeps_timelines_and_confidence_without_request_or_url_metadata(self): original = job(turn(0, 10, confidence={"voice-A": 91}, url="https://secret/segment"), request={"key": "test-secret-key"}, url="https://secret/job") original["output"].update(diarization=[turn(2, 9, "b"), turn(0, 10, "a")], warning="constraints adjusted https://secret/warning", url="https://secret/output") fake = FakeTransport(polls=[original]) with tempfile.TemporaryDirectory() as directory: result, log = self.execute(fake, directory) saved = (Path(directory) / "work" / "result.json").read_text() self.assertEqual(set(result), {"jobId", "status", "output"}) self.assertEqual(set(result["output"]), {"diarization", "exclusiveDiarization"}) self.assertEqual(result["output"]["diarization"], original["output"]["diarization"]) self.assertEqual(result["output"]["exclusiveDiarization"][0]["confidence"], {"voice-A": 91}) for value in ("https://", "test-secret-key", "request", "warning"): self.assertNotIn(value, saved) self.assertIn("constraints adjusted", log) self.assertNotIn("https://", log) def test_automatic_count_is_omitted_and_created_pending_running_are_polled(self): fake = FakeTransport(polls=[{"status": status} for status in ("created", "pending", "running")] + [job(turn(0, 10))]) with tempfile.TemporaryDirectory() as directory: _, log = self.execute(fake, directory) self.assertNotIn("numSpeakers", fake.submissions[0]) self.assertEqual(len(fake.submissions), 1) self.assertEqual(sum("GET" in call for call in fake.calls), 4) self.assertIn("status: created", log) self.assertIn("status: pending", log) self.assertIn("status: running", log) def test_transient_get_http_and_transport_errors_repoll_same_job(self): fake = FakeTransport(polls=[(503, {"error": "try again"}, 0), (0, "", 28), job(turn(0, 10))]) with tempfile.TemporaryDirectory() as directory: _, log = self.execute(fake, directory) urls = [call[-1] for call in fake.calls if "GET" in call] self.assertEqual(urls, ["https://api.pyannote.ai/v1/jobs/test-job"] * 3) self.assertEqual(len(fake.submissions), 1) self.assertEqual(len(fake.uploads), 1) self.assertIn("retrying the same job", log) self.assertNotIn("secret=signature", log) def test_post_and_upload_failures_are_never_retried_and_leave_no_result(self): for endpoint in ("/media/input", "storage.example", "/diarize"): fake = FakeTransport(failure=(endpoint, 503, {"error": "provider unavailable"}, 0)) with self.subTest(endpoint=endpoint), tempfile.TemporaryDirectory() as directory: with self.assertRaisesRegex(ValueError, "HTTP 503"): self.execute(fake, directory) self.assertEqual(list((Path(directory) / "work").iterdir()), []) self.assertEqual(sum(endpoint in call[-1] for call in fake.calls), 1) def test_terminal_and_malformed_poll_failures_never_yield_partial_result(self): failures = [dict(job(), status="failed"), dict(job(), status="canceled"), {"status": "unexpected"}, {"status": "succeeded", "output": {}}, dict(job(turn(0, 10)), jobId="wrong-job"), job(turn(0, 31)), (401, {"error": "bad key"}, 0), (200, "not JSON", 0)] for response in failures: with self.subTest(response=response), tempfile.TemporaryDirectory() as directory: with self.assertRaises(ValueError): self.execute(FakeTransport(polls=[response]), directory) self.assertEqual(list((Path(directory) / "work").iterdir()), []) def test_error_details_redact_key_and_urls(self): fake = FakeTransport(failure=("/media/input", 401, {"error": { "message": "invalid test-secret-key Bearer abc https://storage.example/?token=secret"}}, 0)) with tempfile.TemporaryDirectory() as directory, self.assertRaises(ValueError) as caught: self.execute(fake, directory) self.assertIn("invalid", str(caught.exception)) for secret in ("test-secret-key", "abc", "https://", "token=secret"): self.assertNotIn(secret, str(caught.exception)) def test_missing_key_or_invalid_model_timeout_and_count_prevent_network(self): with tempfile.TemporaryDirectory() as directory, patch.dict(os.environ, {}, clear=True), \ patch.object(diarization.subprocess, "run") as provider: with self.assertRaisesRegex(ValueError, "PYANNOTE_API_KEY"): diarization.run(write_audio(directory), directory) provider.assert_not_called() for kwargs in ({"model": "google/gemini"}, {"timeout": 0}, {"timeout": float("inf")}, {"timeout": True}, {"speakers": 0}): with self.subTest(kwargs=kwargs), tempfile.TemporaryDirectory() as directory: fake = FakeTransport() with self.assertRaises(ValueError): self.execute(fake, directory, **kwargs) self.assertEqual(fake.calls, []) def test_overall_deadline_is_bounded_even_while_waiting(self): with tempfile.TemporaryDirectory() as directory, patch.object(diarization.time, "monotonic", return_value=10): client = diarization.Client(directory, "key", 9) with self.assertRaisesRegex(ValueError, "timed out"): client.request("GET", "https://api.pyannote.ai/v1/jobs/one") with tempfile.TemporaryDirectory() as directory, patch.object(diarization.time, "monotonic", return_value=10), \ patch.object(diarization.time, "sleep") as sleep: diarization.Client(directory, "key", 10.5).pause() sleep.assert_called_once_with(0.5) def test_pending_job_times_out_without_restarting_upload_or_submission(self): fake = FakeTransport(polls=[{"status": "running"}] * 10) clock = [0] with tempfile.TemporaryDirectory() as directory, \ patch.dict(os.environ, {"PYANNOTE_API_KEY": "test-secret-key"}), \ patch.object(diarization.subprocess, "run", side_effect=fake), \ patch.object(diarization.time, "monotonic", side_effect=lambda: clock[0]), \ patch.object(diarization.time, "sleep", side_effect=lambda seconds: clock.__setitem__(0, clock[0] + seconds)), \ contextlib.redirect_stderr(io.StringIO()), self.assertRaisesRegex(ValueError, "timed out"): diarization.run(write_audio(directory), Path(directory) / "work", timeout=5) self.assertEqual(len(fake.submissions), 1) self.assertEqual(len(fake.uploads), 1) self.assertEqual(sum("GET" in call for call in fake.calls), 3) def test_upload_size_limit_prevents_any_http_request(self): fake = FakeTransport() with tempfile.TemporaryDirectory() as directory, patch.object(diarization, "MAX_UPLOAD_BYTES", 2), \ self.assertRaisesRegex(ValueError, "1 GiB"): self.execute(fake, directory) self.assertEqual(len(fake.calls), 1) def test_invalid_upload_url_never_receives_audio_or_bearer(self): for url in ("http://storage.example/upload", "file:///tmp/upload", "https://key@storage.example/upload"): with self.subTest(url=url), tempfile.TemporaryDirectory() as directory, \ patch.object(diarization.subprocess, "run") as transport, self.assertRaises(ValueError): diarization.Client(directory, "key", float("inf")).request("PUT", url, authenticated=False) transport.assert_not_called() if __name__ == "__main__": unittest.main()