backend/SQL_Agent/tests/test_compose_tool_workflow.py

from backend.SQL_Agent.data_sources_sql_toolkit import DataSourcesSQLToolkit


class _MockResponse:
    def __init__(self, status_code, payload):
        self.status_code = status_code
        self._payload = payload
        self.text = str(payload)

    def json(self):
        return self._payload


class _MockClient:
    def post(self, url, json=None, timeout=None, headers=None):
        if url == "/api/v1/tenant-files/datasets/compose":
            return _MockResponse(201, {
                "dataset_id": "dataset_123",
                "status": "materialized",
                "artifact_asset_id": "asset_999",
                "output_format": "parquet",
                "lineage": {"nodes": [], "edges": []}
            })
        return _MockResponse(404, {"detail": "not found"})


def test_compose_dataset_tool_calls_tenant_files_api():
    toolkit = DataSourcesSQLToolkit(api_base_url="http://test")
    toolkit.client = _MockClient()

    class _RunContext:
        session_state = {"tenant_id": "tenant_1", "supabase_jwt": "token"}

    result = toolkit.compose_dataset_workflow(
        run_context=_RunContext(),
        name="model_dataset",
        target_format="parquet",
        steps=[{"step_id": "s1", "type": "sql", "source_name": "postgres_a", "sql_query": "SELECT 1"}],
    )

    assert result["dataset_id"] == "dataset_123"
    assert result["artifact_asset_id"] == "asset_999"