|
| 1 | +"""Tests for table relationship metadata.""" |
| 2 | + |
| 3 | +from pathlib import Path |
| 4 | + |
| 5 | +from agentic_data_contracts.core.contract import DataContract |
| 6 | +from agentic_data_contracts.core.schema import ( |
| 7 | + AllowedTable, |
| 8 | + DataContractSchema, |
| 9 | + SemanticConfig, |
| 10 | +) |
| 11 | +from agentic_data_contracts.semantic.cube import CubeSource |
| 12 | +from agentic_data_contracts.semantic.dbt import DbtSource |
| 13 | +from agentic_data_contracts.semantic.yaml_source import YamlSource |
| 14 | + |
| 15 | + |
| 16 | +def test_yaml_source_loads_relationships(fixtures_dir: Path) -> None: |
| 17 | + source = YamlSource(fixtures_dir / "semantic_source.yml") |
| 18 | + rels = source.get_relationships() |
| 19 | + assert len(rels) == 1 |
| 20 | + assert rels[0].from_ == "analytics.orders.customer_id" |
| 21 | + assert rels[0].to == "analytics.customers.id" |
| 22 | + assert rels[0].type == "many_to_one" |
| 23 | + |
| 24 | + |
| 25 | +def test_yaml_source_no_relationships(tmp_path: Path) -> None: |
| 26 | + (tmp_path / "empty.yml").write_text("metrics: []") |
| 27 | + source = YamlSource(tmp_path / "empty.yml") |
| 28 | + assert source.get_relationships() == [] |
| 29 | + |
| 30 | + |
| 31 | +def test_dbt_source_returns_empty_relationships( |
| 32 | + fixtures_dir: Path, |
| 33 | +) -> None: |
| 34 | + source = DbtSource(fixtures_dir / "sample_dbt_manifest.json") |
| 35 | + assert source.get_relationships() == [] |
| 36 | + |
| 37 | + |
| 38 | +def test_cube_source_returns_empty_relationships( |
| 39 | + fixtures_dir: Path, |
| 40 | +) -> None: |
| 41 | + source = CubeSource(fixtures_dir / "sample_cube_schema.yml") |
| 42 | + assert source.get_relationships() == [] |
| 43 | + |
| 44 | + |
| 45 | +def test_system_prompt_includes_relationships( |
| 46 | + fixtures_dir: Path, |
| 47 | +) -> None: |
| 48 | + source = YamlSource(fixtures_dir / "semantic_source.yml") |
| 49 | + schema = DataContractSchema( |
| 50 | + name="test", |
| 51 | + semantic=SemanticConfig( |
| 52 | + allowed_tables=[ |
| 53 | + AllowedTable.model_validate( |
| 54 | + {"schema": "analytics", "tables": ["orders", "customers"]} |
| 55 | + ), |
| 56 | + ], |
| 57 | + ), |
| 58 | + ) |
| 59 | + dc = DataContract(schema) |
| 60 | + prompt = dc.to_system_prompt(semantic_source=source) |
| 61 | + assert "Table Relationships" in prompt |
| 62 | + assert "analytics.orders.customer_id" in prompt |
| 63 | + assert "analytics.customers.id" in prompt |
| 64 | + assert "many_to_one" in prompt |
| 65 | + |
| 66 | + |
| 67 | +def test_system_prompt_no_relationships_when_empty( |
| 68 | + fixtures_dir: Path, |
| 69 | +) -> None: |
| 70 | + source = DbtSource(fixtures_dir / "sample_dbt_manifest.json") |
| 71 | + schema = DataContractSchema( |
| 72 | + name="test", |
| 73 | + semantic=SemanticConfig( |
| 74 | + allowed_tables=[ |
| 75 | + AllowedTable.model_validate( |
| 76 | + {"schema": "analytics", "tables": ["orders"]} |
| 77 | + ), |
| 78 | + ], |
| 79 | + ), |
| 80 | + ) |
| 81 | + dc = DataContract(schema) |
| 82 | + prompt = dc.to_system_prompt(semantic_source=source) |
| 83 | + assert "Table Relationships" not in prompt |
0 commit comments