110 lines · 3.3 KB
python
| 1 | from __future__ import annotations |
| 2 | |
| 3 | import sys |
| 4 | from pathlib import Path |
| 5 | |
| 6 | sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| 7 | |
| 8 | from space_tape.clip import clip_cues, highlights, label_for, parse_ts, quote_clip, score_cue |
| 9 | |
| 10 | |
| 11 | CUES = [ |
| 12 | {"start": 0.0, "end": 11.78, "speaker": "silence", "text": ""}, |
| 13 | { |
| 14 | "start": 11.78, |
| 15 | "end": 17.46, |
| 16 | "speaker": "DavidSacks", |
| 17 | "text": "I think we're broadcasting. Man, I think we melted the internet there.", |
| 18 | }, |
| 19 | { |
| 20 | "start": 18.66, |
| 21 | "end": 25.7, |
| 22 | "speaker": "elonmusk", |
| 23 | "text": "That was insane, sorry. We are actually doing this from David Sacks' Twitter account because it looks like doing it from mine basically broke the Twitter system.", |
| 24 | }, |
| 25 | { |
| 26 | "start": 51.38, |
| 27 | "end": 52.58, |
| 28 | "speaker": "GovRonDeSantis", |
| 29 | "text": "I'm here.", |
| 30 | }, |
| 31 | { |
| 32 | "start": 125.34, |
| 33 | "end": 131.02, |
| 34 | "speaker": "GovRonDeSantis", |
| 35 | "text": "Well, I am running for president of the United States to lead our great American comeback.", |
| 36 | }, |
| 37 | { |
| 38 | "start": 131.42, |
| 39 | "end": 164.22, |
| 40 | "speaker": "GovRonDeSantis", |
| 41 | "text": "American decline is not inevitable. It is a choice.", |
| 42 | }, |
| 43 | ] |
| 44 | |
| 45 | |
| 46 | def test_parse_ts() -> None: |
| 47 | assert parse_ts(18) == 18.0 |
| 48 | assert parse_ts("18.66") == 18.66 |
| 49 | assert parse_ts("00:18") == 18.0 |
| 50 | assert parse_ts("1:02:03") == 3723.0 |
| 51 | assert parse_ts("00:01:36.2") == 96.2 |
| 52 | |
| 53 | |
| 54 | def test_clip_cues_trims_window() -> None: |
| 55 | out = clip_cues(CUES, 18.0, 26.0) |
| 56 | assert len(out) == 1 |
| 57 | assert out[0]["speaker"] == "elonmusk" |
| 58 | assert out[0]["start"] == 18.66 |
| 59 | assert out[0]["end"] == 25.7 |
| 60 | |
| 61 | |
| 62 | def test_clip_cues_spans_speakers() -> None: |
| 63 | out = clip_cues(CUES, 10.0, 20.0) |
| 64 | voiced = [c["speaker"] for c in out if c["text"]] |
| 65 | assert voiced == ["DavidSacks", "elonmusk"] |
| 66 | assert out[-1]["end"] == 20.0 |
| 67 | |
| 68 | |
| 69 | def test_quote_keeps_timestamp_and_handle() -> None: |
| 70 | md = quote_clip(clip_cues(CUES, 18.0, 26.0), source_url="https://x.com/i/spaces/1eaJbrAlZjjJX") |
| 71 | assert "**@elonmusk**" in md |
| 72 | assert "[00:00:19]" in md or "[00:00:18]" in md |
| 73 | assert "That was insane, sorry." in md |
| 74 | assert "https://x.com/i/spaces/1eaJbrAlZjjJX" in md |
| 75 | |
| 76 | |
| 77 | def test_highlights_pick_the_chops() -> None: |
| 78 | hits = highlights(CUES, max_n=5) |
| 79 | texts = " ".join(h["text"].lower() for h in hits) |
| 80 | assert "melted the internet" in texts |
| 81 | assert "insane" in texts |
| 82 | assert "running for president" in texts |
| 83 | assert all(h["start"] <= hits[i + 1]["start"] for i, h in enumerate(hits[:-1])) |
| 84 | assert all("label" in h for h in hits) |
| 85 | |
| 86 | |
| 87 | def test_silence_does_not_score() -> None: |
| 88 | assert score_cue(CUES[0]) == 0.0 |
| 89 | |
| 90 | |
| 91 | def test_label_truncates() -> None: |
| 92 | assert label_for("One two three") == "One two three" |
| 93 | assert "running for president" in label_for( |
| 94 | "Well, I am running for president of the United States to lead our great American comeback." |
| 95 | ).lower() |
| 96 | assert "melted the internet" in label_for( |
| 97 | "I think we're broadcasting. Man, I think we melted the internet there." |
| 98 | ).lower() |
| 99 | |
| 100 | |
| 101 | if __name__ == "__main__": |
| 102 | test_parse_ts() |
| 103 | test_clip_cues_trims_window() |
| 104 | test_clip_cues_spans_speakers() |
| 105 | test_quote_keeps_timestamp_and_handle() |
| 106 | test_highlights_pick_the_chops() |
| 107 | test_silence_does_not_score() |
| 108 | test_label_truncates() |
| 109 | print("ok") |