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:
@@ -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
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user