Refactor internal data passing to dataclasses and enums, and remove unneeded defensive code

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
2026-08-04 15:13:31 +08:00
co-authored by Claude Fable 5
parent 8d223177cd
commit e0eb343ce0
11 changed files with 540 additions and 333 deletions
+24 -23
View File
@@ -20,6 +20,7 @@ from pop_fem_audit_tools.database import DataSource
from pop_fem_audit_tools.models import (
Artist,
ChartEntry,
Role,
Song,
)
@@ -30,47 +31,47 @@ class TestParseArtistCredit(unittest.TestCase):
def test_plain_solo(self) -> None:
"""Test a plain solo artist credit."""
self.assertEqual(build_db.parse_artist_credit("Adele"),
[("Adele", "primary")])
[("Adele", Role.PRIMARY)])
def test_featuring_with_and(self) -> None:
"""Test a featuring credit with an "and" delimiter."""
self.assertEqual(
build_db.parse_artist_credit(
"Drake featuring Wizkid and Kyla"),
[("Drake", "primary"),
("Wizkid", "featured"),
("Kyla", "featured")])
[("Drake", Role.PRIMARY),
("Wizkid", Role.FEATURED),
("Kyla", Role.FEATURED)])
def test_comma_and_ampersand(self) -> None:
"""Test a credit with comma and ampersand delimiters."""
self.assertEqual(
build_db.parse_artist_credit(
"Lady Gaga, Bradley Cooper & BloodPop"),
[("Lady Gaga", "primary"),
("Bradley Cooper", "primary"),
("BloodPop", "primary")])
[("Lady Gaga", Role.PRIMARY),
("Bradley Cooper", Role.PRIMARY),
("BloodPop", Role.PRIMARY)])
def test_x_delimiter(self) -> None:
"""Test the "x" delimiter."""
self.assertEqual(
build_db.parse_artist_credit("KAROL G x Nicki Minaj"),
[("KAROL G", "primary"),
("Nicki Minaj", "primary")])
[("KAROL G", Role.PRIMARY),
("Nicki Minaj", Role.PRIMARY)])
def test_plus_delimiter(self) -> None:
"""Test the "+" delimiter."""
self.assertEqual(
build_db.parse_artist_credit("Marshmello + Halsey"),
[("Marshmello", "primary"),
("Halsey", "primary")])
[("Marshmello", Role.PRIMARY),
("Halsey", Role.PRIMARY)])
def test_with_delimiter(self) -> None:
"""Test the "with" delimiter."""
self.assertEqual(
build_db.parse_artist_credit(
"Kane Brown with Lauren Alaina"),
[("Kane Brown", "primary"),
("Lauren Alaina", "primary")])
[("Kane Brown", Role.PRIMARY),
("Lauren Alaina", Role.PRIMARY)])
def test_feat_abbreviation(self) -> None:
"""Test that "Feat." splits the featured side."""
@@ -78,17 +79,17 @@ class TestParseArtistCredit(unittest.TestCase):
build_db.parse_artist_credit(
"Ariana Grande Feat. Doja Cat"
" & Megan Thee Stallion"),
[("Ariana Grande", "primary"),
("Doja Cat", "featured"),
("Megan Thee Stallion", "featured")])
[("Ariana Grande", Role.PRIMARY),
("Doja Cat", Role.FEATURED),
("Megan Thee Stallion", Role.FEATURED)])
def test_case_insensitive_featuring(self) -> None:
"""Test that "Featuring" splits case-insensitively."""
self.assertEqual(
build_db.parse_artist_credit(
"24kGoldn Featuring iann dior"),
[("24kGoldn", "primary"),
("iann dior", "featured")])
[("24kGoldn", Role.PRIMARY),
("iann dior", Role.FEATURED)])
class TestBuildDB(unittest.TestCase):
@@ -183,8 +184,8 @@ class TestBuildDB(unittest.TestCase):
[(2016, 2), (2017, 1)])
self.assertEqual([(x.artist.name, x.role, x.position)
for x in song.song_artists],
[("Drake", "primary", 0),
("Wizkid", "featured", 1)])
[("Drake", Role.PRIMARY, 0),
("Wizkid", Role.FEATURED, 1)])
self.assertIn("3 songs", stderr)
self.assertIn("4 chart entries", stderr)
self.assertIn("4 artists", stderr)
@@ -244,11 +245,11 @@ class TestBuildDB(unittest.TestCase):
def test_overrides_apply_over_wikidata(self) -> None:
"""Test that the overrides win over the Wikidata snapshot."""
Path("data/artists_wikidata.csv").write_text(
"name,qid,gender,artist_type,genre,country,note\n"
"name,qid,gender,type,genre,country,note\n"
"Adele,Q2831,female,solo,pop,GB,\n",
encoding="utf-8")
Path("data/artists_overrides.csv").write_text(
"name,qid,gender,artist_type,genre,country,note\n"
"name,qid,gender,type,genre,country,note\n"
"Adele,,,,soul,,manually checked\n",
encoding="utf-8")
self.assertEqual(self.__run_build()[0], 0)
@@ -264,7 +265,7 @@ class TestBuildDB(unittest.TestCase):
def test_unknown_override_name_fails(self) -> None:
"""Test that an unknown override name fails the build."""
Path("data/artists_overrides.csv").write_text(
"name,qid,gender,artist_type,genre,country,note\n"
"name,qid,gender,type,genre,country,note\n"
"Adel,,female,,,,typo\n", encoding="utf-8")
status: int
stderr: str
+1 -1
View File
@@ -26,7 +26,7 @@ class TestFetchArtists(unittest.TestCase):
"""Test cases for the artist metadata fetcher."""
HEADER: list[str] = [
"name", "qid", "gender", "artist_type", "genre",
"name", "qid", "gender", "type", "genre",
"country", "note"]
"""The expected header row of the snapshot CSV file."""
+2 -1
View File
@@ -21,6 +21,7 @@ from pop_fem_audit_tools import config, fetch_lyrics
from pop_fem_audit_tools.database import Base, DataSource
from pop_fem_audit_tools.models import (
Artist,
Role,
Song,
SongArtist,
)
@@ -81,7 +82,7 @@ class TestFetchLyrics(unittest.TestCase):
session.add(song)
session.add(SongArtist(song=song,
artist=artists[artist],
role="primary",
role=Role.PRIMARY,
position=0))
session.commit()
finally:
+5 -4
View File
@@ -14,6 +14,7 @@ from pop_fem_audit_tools.database import Base, DataSource
from pop_fem_audit_tools.models import (
Artist,
ChartEntry,
Role,
Song,
SongArtist,
)
@@ -42,9 +43,9 @@ class TestModels(unittest.TestCase):
song.chart_entries = [ChartEntry(year=2016, rank=4)]
song.song_artists = [
SongArtist(artist=Artist(name="Drake"),
role="primary", position=0),
role=Role.PRIMARY, position=0),
SongArtist(artist=Artist(name="Wizkid"),
role="featured", position=1)]
role=Role.FEATURED, position=1)]
song.lyrics = "Baby, I like your style"
self.__session.add(song)
self.__session.commit()
@@ -63,8 +64,8 @@ class TestModels(unittest.TestCase):
[(2016, 4)])
self.assertEqual([(x.artist.name, x.role, x.position)
for x in song.song_artists],
[("Drake", "primary", 0),
("Wizkid", "featured", 1)])
[("Drake", Role.PRIMARY, 0),
("Wizkid", Role.FEATURED, 1)])
self.assertEqual(song.lyrics,
"Baby, I like your style")
artist: Artist | None = self.__session.scalar(
+43 -33
View File
@@ -51,8 +51,12 @@ class RunLLMTestCase(unittest.TestCase):
error_type: str) -> mock.Mock:
"""Create a mock errored batch result entry.
The error object is shaped as the SDK envelope: the outer
error carries the constant type ``error``, and the specific
error code lives at ``error.error.type``.
:param custom_id: The custom ID of the entry.
:param error_type: The error type.
:param error_type: The specific error code.
:return: The mock result entry.
"""
entry: mock.Mock = mock.Mock()
@@ -60,7 +64,9 @@ class RunLLMTestCase(unittest.TestCase):
entry.result = mock.Mock()
entry.result.type = "errored"
entry.result.error = mock.Mock()
entry.result.error.type = error_type
entry.result.error.type = "error"
entry.result.error.error = mock.Mock()
entry.result.error.error.type = error_type
return entry
@staticmethod
@@ -105,9 +111,10 @@ class TestLoadItems(RunLLMTestCase):
path: Path = self.__write_input(
'{"id": "a", "content": "one"}\n'
'{"id": "b", "content": "two"}\n')
items: list[run_llm.Item] = run_llm.load_items(path)
self.assertEqual(items, [{"id": "a", "content": "one"},
{"id": "b", "content": "two"}])
items: list[run_llm.InputItem] = run_llm.load_items(path)
self.assertEqual(items, [
run_llm.InputItem(id="a", content="one"),
run_llm.InputItem(id="b", content="two")])
def test_malformed_json_names_line(self) -> None:
"""Test that malformed JSON reports the line number."""
@@ -161,7 +168,7 @@ class TestRequestBuilding(RunLLMTestCase):
def test_build_request(self) -> None:
"""Test the shape of a batch request."""
request: dict[str, Any] = run_llm.build_request(
{"id": "song-1", "content": "the lyrics"},
run_llm.InputItem(id="song-1", content="the lyrics"),
"the system prompt", 2048)
self.assertEqual(request["custom_id"], "song-1")
params: dict[str, Any] = request["params"]
@@ -188,18 +195,18 @@ class TestAgreement(RunLLMTestCase):
def test_split_by_agreement(self) -> None:
"""Test splitting items into agreed and disagreeing ones."""
items: list[run_llm.Item] = [
{"id": "a", "content": "one"},
{"id": "b", "content": "two"},
{"id": "c", "content": "three"}]
items: list[run_llm.InputItem] = [
run_llm.InputItem(id="a", content="one"),
run_llm.InputItem(id="b", content="two"),
run_llm.InputItem(id="c", content="three")]
run1: run_llm.Results = {
"a": {"id": "a", "text": "same\n"},
"b": {"id": "b", "text": "left"},
"c": {"id": "c", "text": " padded "}}
"a": run_llm.BatchResult(id="a", text="same\n"),
"b": run_llm.BatchResult(id="b", text="left"),
"c": run_llm.BatchResult(id="c", text=" padded ")}
run2: run_llm.Results = {
"a": {"id": "a", "text": "same"},
"b": {"id": "b", "text": "right"},
"c": {"id": "c", "text": "padded"}}
"a": run_llm.BatchResult(id="a", text="same"),
"b": run_llm.BatchResult(id="b", text="right"),
"c": run_llm.BatchResult(id="c", text="padded")}
agreed: list[str]
disagreed: list[str]
agreed, disagreed = run_llm.split_by_agreement(
@@ -213,14 +220,15 @@ class TestFinalRecords(RunLLMTestCase):
def test_build_final_records(self) -> None:
"""Test assembling final records from runs and arbitration."""
items: list[run_llm.Item] = [
{"id": "a", "content": "one"},
{"id": "b", "content": "two"}]
items: list[run_llm.InputItem] = [
run_llm.InputItem(id="a", content="one"),
run_llm.InputItem(id="b", content="two")]
run1: run_llm.Results = {
"a": {"id": "a", "text": "agreed text\n"},
"b": {"id": "b", "text": "left"}}
"a": run_llm.BatchResult(id="a", text="agreed text\n"),
"b": run_llm.BatchResult(id="b", text="left")}
arbitration: run_llm.Results = {
"b": {"id": "b", "text": "arbitrated text"}}
"b": run_llm.BatchResult(id="b",
text="arbitrated text")}
records: list[dict[str, str]] = run_llm.build_final_records(
items, run1, arbitration)
self.assertEqual(records, [
@@ -237,23 +245,23 @@ class TestCollectResults(RunLLMTestCase):
client: mock.Mock = mock.Mock()
client.messages.batches.results.return_value = iter([
self._make_success_entry("a", "output a"),
self._make_error_entry("b", "invalid_request")])
self._make_error_entry("b", "invalid_request_error")])
results: run_llm.Results = run_llm.collect_results(
client, "batch_x")
self.assertEqual(results["a"]["text"], "output a")
self.assertEqual(results["a"]["stop_reason"], "end_turn")
self.assertEqual(results["a"]["usage"],
self.assertEqual(results["a"].text, "output a")
self.assertEqual(results["a"].stop_reason, "end_turn")
self.assertEqual(results["a"].usage,
{"input_tokens": 10, "output_tokens": 5})
self.assertEqual(results["b"],
{"id": "b", "error": "invalid_request"})
self.assertEqual(results["b"], run_llm.BatchResult(
id="b", error="invalid_request_error"))
client.messages.batches.results.assert_called_once_with(
"batch_x")
def test_find_failures(self) -> None:
"""Test finding failed and missing items."""
results: run_llm.Results = {
"a": {"id": "a", "text": "fine"},
"b": {"id": "b", "error": "errored"}}
"a": run_llm.BatchResult(id="a", text="fine"),
"b": run_llm.BatchResult(id="b", error="errored")}
self.assertEqual(
run_llm.find_failures(["a", "b", "c"], results),
["b", "c"])
@@ -345,7 +353,7 @@ class TestMainFlow(RunLLMTestCase):
mock.Mock(id=x) for x in batch_ids]
ended: mock.Mock = mock.Mock()
ended.processing_status = "ended"
ended.ended_at = "2026-07-30T20:00:00+08:00"
ended.ended_at = datetime(2026, 7, 30, 20, 0)
client.messages.batches.retrieve.return_value = ended
results: dict[str, list[Any]] = {
"batch_run1": run1, "batch_run2": run2,
@@ -482,7 +490,8 @@ class TestMainFlow(RunLLMTestCase):
"""Test that a failed item aborts with a non-zero status."""
client: mock.Mock = self.__make_client(
run1=[self._make_success_entry("a", "answer a"),
self._make_error_entry("b", "invalid_request")],
self._make_error_entry(
"b", "invalid_request_error")],
run2=[self._make_success_entry("a", "answer a"),
self._make_success_entry("b", "answer b")])
status: int = self.__run_main(self.__argv, client)[0]
@@ -495,7 +504,8 @@ class TestMainFlow(RunLLMTestCase):
run1_lines: list[str] = (run_dir / "run1.jsonl") \
.read_text(encoding="utf-8").splitlines()
self.assertEqual(json.loads(run1_lines[1]),
{"id": "b", "error": "invalid_request"})
{"id": "b",
"error": "invalid_request_error"})
def test_invalid_input_exits_non_zero(self) -> None:
"""Test that an invalid input file aborts before archiving."""