DataOps-platform/tests/data_research/test_structured_extractors.py

from __future__ import annotations

import io

import pandas as pd
import pytest


def test_sql_extractor_returns_multi_table_candidates_with_statement_evidence():
    from app.core.data_research.extractors.base import ExtractionContext
    from app.core.data_research.extractors.sql import SqlExtractor

    content = b"""
    CREATE TABLE customer (id INTEGER PRIMARY KEY, name VARCHAR(50));
    CREATE TABLE orders (id INTEGER PRIMARY KEY, customer_id INTEGER NOT NULL);
    """
    batch = SqlExtractor().extract(
        content,
        ExtractionContext(filename="schema.sql", media_type="application/sql"),
    )

    assert batch.parser_version == "sql-ddl-v1"
    assert len(batch.content_hash) == 64
    assert [item.data["table_info"]["name_en"] for item in batch.items] == [
        "customer",
        "orders",
    ]
    assert batch.items[0].evidence[0].locator == {
        "kind": "sql.statement",
        "statement": 1,
    }


def test_csv_extractor_detects_bom_delimiter_and_column_types():
    from app.core.data_research.extractors.base import ExtractionContext
    from app.core.data_research.extractors.csv import CsvExtractor

    content = "\ufeffid;amount;active\n1;12.5;true\n2;9.0;false\n".encode("utf-8")
    batch = CsvExtractor().extract(
        content,
        ExtractionContext(filename="orders.csv", media_type="text/csv"),
    )

    table = batch.items[0]
    assert table.data["table_info"]["name_en"] == "orders"
    assert [(column["name_en"], column["data_type"]) for column in table.data["columns"]] == [
        ("id", "integer"),
        ("amount", "decimal"),
        ("active", "boolean"),
    ]
    assert table.data["columns"][1]["evidence_locator"] == {
        "kind": "csv.header",
        "column": 2,
        "name": "amount",
    }


def test_csv_header_only_is_valid_and_empty_file_is_rejected():
    from app.core.data_research.extractors.base import ExtractionContext
    from app.core.data_research.extractors.csv import CsvExtractor

    context = ExtractionContext(filename="empty.csv", media_type="text/csv")
    batch = CsvExtractor().extract(b"id,name\n", context)
    assert [column["data_type"] for column in batch.items[0].data["columns"]] == [
        "string",
        "string",
    ]

    with pytest.raises(ValueError, match="header"):
        CsvExtractor().extract(b"", context)


def test_excel_extractor_returns_one_candidate_per_sheet_with_cell_evidence():
    from app.core.data_research.extractors.base import ExtractionContext
    from app.core.data_research.extractors.excel import ExcelExtractor

    stream = io.BytesIO()
    with pd.ExcelWriter(stream, engine="openpyxl") as writer:
        pd.DataFrame({"customer_id": [1], "name": ["Ada"]}).to_excel(
            writer, sheet_name="Customers", index=False
        )
        pd.DataFrame({"order_id": [9], "total": [3.5]}).to_excel(
            writer, sheet_name="Orders", index=False
        )

    batch = ExcelExtractor().extract(
        stream.getvalue(),
        ExtractionContext(
            filename="dictionary.xlsx",
            media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        ),
    )

    assert [item.data["table_info"]["name_en"] for item in batch.items] == [
        "Customers",
        "Orders",
    ]
    assert batch.items[1].data["columns"][1]["evidence_locator"] == {
        "kind": "excel.cell_range",
        "sheet": "Orders",
        "range": "B1:B2",
    }


def test_registry_selects_by_filename_and_enforces_input_limit():
    from app.core.data_research.extractors.base import ExtractionContext
    from app.core.data_research.extractors.registry import default_registry

    registry = default_registry()
    assert registry.resolve("text/csv", "orders.csv").parser_version == "csv-v1"
    assert registry.resolve("application/sql", "schema.sql").parser_version == "sql-ddl-v1"

    context = ExtractionContext(
        filename="orders.csv",
        media_type="text/csv",
        max_bytes=4,
    )
    with pytest.raises(ValueError, match="size limit"):
        registry.extract(b"id,name\n1,Ada\n", context)