from gti_discovery.company_context import collect_company_context


class Reader:
    def company_context_rows(self):
        return [
            {"UF": "SP", "REGIME": "SIMPLES NACIONAL", "CNPJ": "12.345.678/0001-90"},
            {"UF": "SP", "REGIME": "SIMPLES NACIONAL", "CNPJ": "12.345.678/0001-90"},
        ]


def test_company_context_is_targeted_and_redacts_cnpj():
    result = collect_company_context(Reader())

    assert result == {"uf": "SP", "regime": "SIMPLES NACIONAL", "cnpj": "REDACTED", "record_count": 2}


def test_company_context_rejects_generic_readers():
    class Generic:
        def execute(self, sql):
            return []

    try:
        collect_company_context(Generic())
    except TypeError as error:
        assert "company_context_rows" in str(error)
    else:
        raise AssertionError("generic query access must be rejected")