tests/agents/d3/test_extractor.py

from app.agents.d3.extractor import D3DataExtractor
from app.agents.d3.registry import TEMPLATES
from app.agents.d3 import templates  # noqa: F401
from app.agents.d3.templates._hierarchy import HierNode, HierarchyData


class _FakeClient:
    def __init__(self, obj): self._obj = obj
    def structured_complete(self, messages, schema, **kw): return self._obj


def test_fill_returns_schema_instance():
    t = TEMPLATES["bar_chart"]
    filled = t.schema(title="GDP", x_label="Country", y_label="T$", data=[{"name": "US", "value": 21}])
    ex = D3DataExtractor(client=_FakeClient(filled))
    out = ex.fill(t, "GDP", [{"source": "s", "text": "US 21 trillion"}], "graduate")
    assert out is not None and out.data[0].name == "US"


def test_fill_never_declines_when_model_returns_empty_data():
    # D3DataExtractor.fill() no longer declines on an empty schema instance -- the
    # extraction prompt now asks the model to synthesize plausible illustrative data
    # itself rather than leave the schema empty, so fill() always returns what the
    # model gave back, whatever its shape.
    t = TEMPLATES["bar_chart"]
    empty = t.schema(data=[])
    ex = D3DataExtractor(client=_FakeClient(empty))
    out = ex.fill(t, "x", [{"source": "s", "text": "no numbers here"}], "graduate")
    assert out is empty


def test_fill_never_declines_hierarchy_root_with_no_children_and_no_value():
    empty_root = HierNode(name="Root")
    filled = HierarchyData(title="Empty", root=empty_root)
    ex = D3DataExtractor(client=_FakeClient(filled))
    out = ex.fill(
        TEMPLATES["bar_chart"],
        "x",
        [{"source": "s", "text": "no hierarchy here"}],
        "graduate",
    )
    assert out is filled


def test_fill_accepts_hierarchy_root_with_value_as_valid():
    # A single-node root WITH a value (or with children) must NOT be declined —
    # per the plan, decline only when root.children == [] AND root.value is None.
    valued_root = HierNode(name="Root", value=42)
    filled = HierarchyData(title="Single node", root=valued_root)
    ex = D3DataExtractor(client=_FakeClient(filled))
    out = ex.fill(
        TEMPLATES["bar_chart"],
        "x",
        [{"source": "s", "text": "root has a value"}],
        "graduate",
    )
    assert out is not None
    assert out.root.value == 42