Build the arbitration input with a per-ID extras merge

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

tools/tests/test_export_llm_input.py

@@ -76,16 +76,21 @@ class TestExportLlmInput(unittest.TestCase):
             session.close()
 
     def __run_export(
-            self, extras: Path | None = None) -> tuple[int, str]:
+            self, extras: Path | None = None,
+            extras_per_id: Path | None = None) -> tuple[int, str]:
         """Run the exporter with the standard error captured.
 
         :param extras: The extras JSON file, or None for none.
+        :param extras_per_id: The per-ID extras JSON file, or None
+            for none.
         :return: A tuple of the exit status and the standard
             error.
         """
         argv: list[str] = [str(self.__output)]
         if extras is not None:
             argv += ["--extras", str(extras)]
+        if extras_per_id is not None:
+            argv += ["--extras-per-id", str(extras_per_id)]
         stderr: io.StringIO = io.StringIO()
         with redirect_stderr(stderr):
             status: int = export_llm_input.main(argv)
@@ -101,6 +106,16 @@ class TestExportLlmInput(unittest.TestCase):
         path.write_text(text, encoding="utf-8")
         return path
 
+    def __write_extras_per_id(self, text: str) -> Path:
+        """Write a per-ID extras file with the given raw text.
+
+        :param text: The raw file content.
+        :return: The path of the written per-ID extras file.
+        """
+        path: Path = self.__dir / "extras-per-id.json"
+        path.write_text(text, encoding="utf-8")
+        return path
+
     @staticmethod
     def __read_records(path: Path) -> list[dict[str, str]]:
         """Read the JSONL records of a file.
@@ -228,3 +243,115 @@ class TestExportLlmInput(unittest.TestCase):
         self.assertEqual(status, 1)
         self.assertIn("error:", stderr)
         self.assertFalse(self.__output.exists())
+
+    def test_extras_per_id_merges_each_song_its_own(self) -> None:
+        """Test that with per-ID extras, each record's content is
+        a JSON object of the lyrics followed by that song's own
+        keys in their file order."""
+        self.__seed([
+            ("Hello", "Adele", "hello lyrics\n"),
+            ("Umbrella", "Rihanna", "umbrella lyrics\n")])
+        per_id: Path = self.__write_extras_per_id(
+            '{"song-1": {"b": 2, "a": 1},'
+            ' "song-2": {"disagreements": {"k": ["q"]}}}')
+        status: int
+        stderr: str
+        status, stderr = self.__run_export(None, per_id)
+        self.assertEqual(status, 0)
+        records: list[dict[str, str]] = self.__read_records(
+            self.__output)
+        self.assertEqual(len(records), 2)
+        first: dict[str, Any] = json.loads(records[0]["content"])
+        self.assertEqual(
+            list(first.keys()), ["lyrics", "b", "a"])
+        self.assertEqual(first["lyrics"], "hello lyrics\n")
+        self.assertEqual(first["b"], 2)
+        self.assertEqual(first["a"], 1)
+        second: dict[str, Any] = json.loads(records[1]["content"])
+        self.assertEqual(
+            list(second.keys()), ["lyrics", "disagreements"])
+        self.assertEqual(second["lyrics"], "umbrella lyrics\n")
+        self.assertEqual(
+            second["disagreements"], {"k": ["q"]})
+
+    def test_extras_per_id_restricts_the_export(self) -> None:
+        """Test that the songs the per-ID extras do not name are
+        not exported."""
+        self.__seed([
+            ("Hello", "Adele", "hello lyrics\n"),
+            ("Umbrella", "Rihanna", "umbrella lyrics\n"),
+            ("Halo", "Beyonce", "halo lyrics\n")])
+        per_id: Path = self.__write_extras_per_id(
+            '{"song-3": {"a": 1}, "song-1": {"a": 2}}')
+        status: int
+        stderr: str
+        status, stderr = self.__run_export(None, per_id)
+        self.assertEqual(status, 0)
+        records: list[dict[str, str]] = self.__read_records(
+            self.__output)
+        self.assertEqual(
+            [x["id"] for x in records], ["song-1", "song-3"])
+        self.assertIn("done: 2 songs exported", stderr)
+
+    def test_extras_per_id_unknown_id_fails(self) -> None:
+        """Test that a per-ID extras file naming a song the
+        working store does not have is rejected."""
+        self.__seed([("Hello", "Adele", "hello lyrics\n")])
+        per_id: Path = self.__write_extras_per_id(
+            '{"song-1": {"a": 1}, "song-9": {"a": 2}}')
+        status: int
+        stderr: str
+        status, stderr = self.__run_export(None, per_id)
+        self.assertEqual(status, 1)
+        self.assertIn("song-9", stderr)
+        self.assertFalse(self.__output.exists())
+
+    def test_extras_per_id_after_the_shared_extras(self) -> None:
+        """Test that with both extras options, a record's keys are
+        the lyrics, the shared extras, then that song's own."""
+        self.__seed([
+            ("Hello", "Adele", "hello lyrics\n"),
+            ("Umbrella", "Rihanna", "umbrella lyrics\n")])
+        extras: Path = self.__write_extras('{"shared": "s"}')
+        per_id: Path = self.__write_extras_per_id(
+            '{"song-2": {"own": "o"}}')
+        status: int
+        stderr: str
+        status, stderr = self.__run_export(extras, per_id)
+        self.assertEqual(status, 0)
+        records: list[dict[str, str]] = self.__read_records(
+            self.__output)
+        self.assertEqual([x["id"] for x in records], ["song-2"])
+        content: dict[str, Any] = json.loads(
+            records[0]["content"])
+        self.assertEqual(
+            list(content.keys()), ["lyrics", "shared", "own"])
+        self.assertEqual(content["lyrics"], "umbrella lyrics\n")
+        self.assertEqual(content["shared"], "s")
+        self.assertEqual(content["own"], "o")
+
+    def test_extras_per_id_non_object_value_fails(self) -> None:
+        """Test that a per-ID extras file whose song does not have
+        a JSON object is rejected."""
+        self.__seed([("Hello", "Adele", "hello lyrics\n")])
+        per_id: Path = self.__write_extras_per_id(
+            '{"song-1": [1, 2]}')
+        status: int
+        stderr: str
+        status, stderr = self.__run_export(None, per_id)
+        self.assertEqual(status, 1)
+        self.assertIn("error:", stderr)
+        self.assertFalse(self.__output.exists())
+
+    def test_extras_per_id_with_lyrics_key_fails(self) -> None:
+        """Test that a per-ID extras file whose song carries a
+        "lyrics" key is rejected."""
+        self.__seed([("Hello", "Adele", "hello lyrics\n")])
+        per_id: Path = self.__write_extras_per_id(
+            '{"song-1": {"lyrics": "not allowed"}}')
+        status: int
+        stderr: str
+        status, stderr = self.__run_export(None, per_id)
+        self.assertEqual(status, 1)
+        self.assertIn("error:", stderr)
+        self.assertFalse(self.__output.exists())