34 lines
1.2 KiB
Python
34 lines
1.2 KiB
Python
|
"""Tests for diarization service data structures and payload conversion."""
|
||
|
|
|
||
|
|
from voice_to_notes.services.diarize import (
|
||
|
|
DiarizationResult,
|
||
|
|
SpeakerSegment,
|
||
|
|
diarization_to_payload,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_diarization_to_payload():
|
||
|
|
result = DiarizationResult(
|
||
|
|
speaker_segments=[
|
||
|
|
SpeakerSegment(speaker="SPEAKER_00", start_ms=0, end_ms=5000),
|
||
|
|
SpeakerSegment(speaker="SPEAKER_01", start_ms=5000, end_ms=10000),
|
||
|
|
SpeakerSegment(speaker="SPEAKER_00", start_ms=10000, end_ms=15000),
|
||
|
|
],
|
||
|
|
num_speakers=2,
|
||
|
|
speakers=["SPEAKER_00", "SPEAKER_01"],
|
||
|
|
)
|
||
|
|
payload = diarization_to_payload(result)
|
||
|
|
assert payload["num_speakers"] == 2
|
||
|
|
assert len(payload["speaker_segments"]) == 3
|
||
|
|
assert payload["speakers"] == ["SPEAKER_00", "SPEAKER_01"]
|
||
|
|
assert payload["speaker_segments"][0]["speaker"] == "SPEAKER_00"
|
||
|
|
assert payload["speaker_segments"][1]["start_ms"] == 5000
|
||
|
|
|
||
|
|
|
||
|
|
def test_diarization_to_payload_empty():
|
||
|
|
result = DiarizationResult()
|
||
|
|
payload = diarization_to_payload(result)
|
||
|
|
assert payload["num_speakers"] == 0
|
||
|
|
assert payload["speaker_segments"] == []
|
||
|
|
assert payload["speakers"] == []
|