tests/test_parsing.py

from __future__ import annotations

import unittest

from src.parsing import extract_finding_discovery_payload, extract_json_payload


REAL_MEDGEMMA_DISCOVERY_RESPONSE = """FINDINGS:
There is a small nodule in the right upper lobe.
There is a small nodule in the left upper lobe.
The heart size is normal.
The mediastinum is normal.
The lungs are clear.
There is no pleural effusion.
There is no pneumothorax.
There is no fracture.

Consolidated list of findings:
[
  {"finding": "nodule", "anatomical_location": "right upper lobe", "certainty": "positive"},
  {"finding": "nodule", "anatomical_location": "left upper lobe", "certainty": "positive"}
]"""


class TestExtractFindingDiscoveryPayload(unittest.TestCase):
    def test_a_strict_expected_dict(self) -> None:
        response = (
            '{"findings": [{"finding": "nodule", "anatomical_location": "right upper lobe", '
            '"certainty": "positive"}]}'
        )
        payload = extract_finding_discovery_payload(response)
        self.assertEqual(len(payload["findings"]), 1)
        self.assertEqual(payload["findings"][0]["finding"], "nodule")

    def test_b_top_level_list(self) -> None:
        response = (
            '[{"finding": "nodule", "anatomical_location": "right upper lobe", '
            '"certainty": "positive"}]'
        )
        payload = extract_finding_discovery_payload(response)
        self.assertEqual(payload, {
            "findings": [
                {
                    "finding": "nodule",
                    "anatomical_location": "right upper lobe",
                    "certainty": "positive",
                }
            ]
        })

    def test_c_exact_real_medgemma_response(self) -> None:
        payload = extract_finding_discovery_payload(REAL_MEDGEMMA_DISCOVERY_RESPONSE)
        self.assertEqual(len(payload["findings"]), 2)
        self.assertEqual(
            payload["findings"][0]["anatomical_location"],
            "right upper lobe",
        )
        self.assertEqual(
            payload["findings"][1]["anatomical_location"],
            "left upper lobe",
        )

    def test_d_fenced_json(self) -> None:
        response = """Here is the result:
```json
{
  "findings": [
    {
      "finding": "pneumothorax",
      "anatomical_location": "right apex",
      "certainty": "positive"
    }
  ]
}
```"""
        payload = extract_finding_discovery_payload(response)
        self.assertEqual(len(payload["findings"]), 1)
        self.assertEqual(payload["findings"][0]["finding"], "pneumothorax")

    def test_e_explicit_empty_findings(self) -> None:
        payload = extract_finding_discovery_payload('{"findings":[]}')
        self.assertEqual(payload, {"findings": []})

    def test_f_top_level_empty_list(self) -> None:
        payload = extract_finding_discovery_payload("[]")
        self.assertEqual(payload, {"findings": []})

    def test_g_malformed_response_raises(self) -> None:
        response = """FINDINGS:
The right upper lobe may contain a nodule."""
        with self.assertRaises(ValueError):
            extract_finding_discovery_payload(response)

    def test_h_multiple_standalone_finding_objects_raises(self) -> None:
        response = """First:
{"finding":"nodule","anatomical_location":"right upper lobe","certainty":"positive"}

Second:
{"finding":"nodule","anatomical_location":"left upper lobe","certainty":"positive"}"""
        with self.assertRaises(ValueError):
            extract_finding_discovery_payload(response)


class TestExtractJsonPayload(unittest.TestCase):
    def test_parses_localization_list_with_trailing_text(self) -> None:
        response = (
            '```json\n'
            '[{"box_2d": [100, 200, 150, 250], "label": "nodule"}]\n'
            '```\n'
            'Localized successfully.'
        )
        payload = extract_json_payload(response, list)
        self.assertEqual(len(payload), 1)


if __name__ == "__main__":
    unittest.main()