feat: support table function fixtures in model tests (#740)

Co-authored-by: Kevin Longe <273747162+kvlonge@users.noreply.github.com>

Author: kvlonge

src/sqlbuild/cli/commands/_helpers/test/inspection.py

@@ -19,6 +19,11 @@ def write_test_plan_inspection(*, stream: TextIO, plan_output: PlanOutput) -> in
         _write_names(stream=stream, label="mocked refs", names=entry.mock_ref_names)
         _write_names(stream=stream, label="mocked seeds", names=entry.mock_seed_names)
         _write_names(stream=stream, label="mocked dbt refs", names=entry.mock_dbt_ref_names)
+        _write_names(
+            stream=stream,
+            label="mocked table functions",
+            names=entry.mock_table_function_names,
+        )
         _write_names(
             stream=stream,
             label="real models",

src/sqlbuild/compiler/compile/_helpers/assembly/project.py

@@ -1439,6 +1439,7 @@ def _assemble_compiled_sql_test(
             mock_source_names=model_payload.mock_source_names,
             mock_seed_names=model_payload.mock_seed_names,
             mock_dbt_ref_names=model_payload.mock_dbt_ref_names,
+            mock_table_function_names=model_payload.mock_table_function_names,
             expected_ctes=model_payload.expected_ctes,
             expected_model_names=model_payload.expected_model_names,
             assertion_ctes=model_payload.assertion_ctes,

src/sqlbuild/compiler/compile/_helpers/attachment/sql_tests.py

@@ -273,6 +273,7 @@ def build_test_inputs(
                     known_model_names=known_model_names,
                     known_seed_names=known_seed_names,
                     known_source_names=known_source_names,
+                    known_table_function_names=known_table_function_names,
                     loaded_macros=loaded_macros,
                 )
                 test_payload: (
@@ -387,6 +388,7 @@ def _build_test_input_payload(
                 mock_source_names=model_payload.mock_source_names,
                 mock_seed_names=model_payload.mock_seed_names,
                 mock_dbt_ref_names=model_payload.mock_dbt_ref_names,
+                mock_table_function_names=model_payload.mock_table_function_names,
                 expected_ctes=model_payload.expected_ctes,
                 expected_model_names=model_payload.expected_model_names,
                 assertion_ctes=model_payload.assertion_ctes,
@@ -678,6 +680,7 @@ def validate_test_ctes(
     known_model_names: set[str],
     known_seed_names: set[str],
     known_source_names: set[str],
+    known_table_function_names: set[str],
     loaded_macros: dict[str, LoadedMacro],
 ) -> None:
     """Validate SQL-native test CTE targets against discovered inputs."""
@@ -705,6 +708,13 @@ def validate_test_ctes(
             raise CompileInputError(
                 f"SQL test file {test_file.relative_path} mocks unknown seed '{mock_seed_name}'"
             )
+    mock_table_function_name: str
+    for mock_table_function_name in model_payload.mock_table_function_names:
+        if mock_table_function_name not in known_table_function_names:
+            raise CompileInputError(
+                f"SQL test file {test_file.relative_path} mocks unknown table function "
+                f"'{mock_table_function_name}'"
+            )
     macro_mock_name: str
     for macro_mock_name in model_payload.macro_mocks:
         if macro_mock_name not in loaded_macros:

src/sqlbuild/compiler/compile/_helpers/sql_tests/core.py

@@ -36,6 +36,7 @@
     SQL_WITH_KEYWORD,
     TABLE_FN_ACTUAL_TEST_CTE_NAME,
     TABLE_FN_EXPECTED_TEST_CTE_NAME,
+    TABLE_FN_TEST_CTE_PREFIX,
     UDF_ACTUAL_TEST_CTE_NAME,
     UDF_EXPECTED_TEST_CTE_NAME,
 )
@@ -392,6 +393,7 @@ def _classify_model_sql_test_ctes(
     mock_source_names: list[str] = []
     mock_seed_names: list[str] = []
     mock_dbt_ref_names: list[str] = []
+    mock_table_function_names: list[str] = []
     expected_ctes: list[CompileSqlTestCte] = []
     expected_model_names: list[str] = []
     assertion_ctes: list[CompileSqlTestCte] = []
@@ -467,6 +469,17 @@ def _classify_model_sql_test_ctes(
             )
             authored_ctes.append(cte)
             continue
+        if cte.name.startswith(TABLE_FN_TEST_CTE_PREFIX):
+            mock_table_function_names.append(
+                _require_prefixed_name(
+                    cte_name=cte.name,
+                    prefix=TABLE_FN_TEST_CTE_PREFIX,
+                    label="__table_fn__<function>",
+                    file_label=file_label,
+                )
+            )
+            authored_ctes.append(cte)
+            continue
         if cte.name.startswith(EXPECTED_TEST_CTE_PREFIX):
             expected_model_names.append(
                 _require_prefixed_name(
@@ -501,11 +514,12 @@ def _classify_model_sql_test_ctes(
         *mock_source_names,
         *mock_seed_names,
         *mock_dbt_ref_names,
+        *mock_table_function_names,
     )
     if not mock_target_names:
         raise CompileInputError(
             f"SQL test '{file_label}' must define at least one __ref__*, __source__*, "
-            "__seed__*, or __dbt_ref__* mock CTE"
+            "__seed__*, __dbt_ref__*, or __table_fn__* mock CTE"
         )
     check_names: tuple[str, ...] = (*expected_model_names, *assertion_names)
     if not check_names:
@@ -520,6 +534,7 @@ def _classify_model_sql_test_ctes(
         mock_source_names=tuple(mock_source_names),
         mock_seed_names=tuple(mock_seed_names),
         mock_dbt_ref_names=tuple(mock_dbt_ref_names),
+        mock_table_function_names=tuple(mock_table_function_names),
         expected_ctes=tuple(expected_ctes),
         expected_model_names=tuple(expected_model_names),
         assertion_ctes=tuple(assertion_ctes),
@@ -536,6 +551,7 @@ def _is_model_mode_cte(cte_name: str) -> bool:
             SOURCE_TEST_CTE_PREFIX,
             SEED_TEST_CTE_PREFIX,
             DBT_REF_TEST_CTE_PREFIX,
+            TABLE_FN_TEST_CTE_PREFIX,
             EXPECTED_TEST_CTE_PREFIX,
             ASSERT_TEST_CTE_PREFIX,
         )

src/sqlbuild/compiler/compile/constants.py

@@ -111,6 +111,7 @@
 SOURCE_TEST_CTE_PREFIX: str = SqlReferenceKind.SOURCE.fixture_cte_prefix
 SEED_TEST_CTE_PREFIX: str = SqlReferenceKind.SEED.fixture_cte_prefix
 DBT_REF_TEST_CTE_PREFIX: str = SqlReferenceKind.DBT_REF.fixture_cte_prefix
+TABLE_FN_TEST_CTE_PREFIX: str = SqlReferenceKind.TABLE_FUNCTION.fixture_cte_prefix
 MACRO_TEST_CTE_PREFIX: str = "__macro__"
 MACRO_ACTUAL_TEST_CTE_NAME: str = "__macro_actual__"
 MACRO_EXPECTED_TEST_CTE_NAME: str = "__macro_expected__"

src/sqlbuild/compiler/compile/models.py

@@ -926,6 +926,7 @@ class CompileModelSqlTestCtes:
     mock_source_names: tuple[str, ...] = field(default_factory=tuple)
     mock_seed_names: tuple[str, ...] = field(default_factory=tuple)
     mock_dbt_ref_names: tuple[str, ...] = field(default_factory=tuple)
+    mock_table_function_names: tuple[str, ...] = field(default_factory=tuple)
     expected_ctes: tuple[CompileSqlTestCte, ...] = field(default_factory=tuple)
     expected_model_names: tuple[str, ...] = field(default_factory=tuple)
     assertion_ctes: tuple[CompileSqlTestCte, ...] = field(default_factory=tuple)
@@ -960,6 +961,7 @@ class CompileModelSqlTestInputPayload:
     mock_source_names: tuple[str, ...] = field(default_factory=tuple)
     mock_seed_names: tuple[str, ...] = field(default_factory=tuple)
     mock_dbt_ref_names: tuple[str, ...] = field(default_factory=tuple)
+    mock_table_function_names: tuple[str, ...] = field(default_factory=tuple)
     expected_ctes: tuple[CompileSqlTestCte, ...] = field(default_factory=tuple)
     expected_model_names: tuple[str, ...] = field(default_factory=tuple)
     assertion_ctes: tuple[CompileSqlTestCte, ...] = field(default_factory=tuple)
@@ -1007,6 +1009,7 @@ class CompiledModelSqlTestPayload:
     mock_source_names: tuple[str, ...] = field(default_factory=tuple)
     mock_seed_names: tuple[str, ...] = field(default_factory=tuple)
     mock_dbt_ref_names: tuple[str, ...] = field(default_factory=tuple)
+    mock_table_function_names: tuple[str, ...] = field(default_factory=tuple)
     expected_ctes: tuple[CompileSqlTestCte, ...] = field(default_factory=tuple)
     expected_model_names: tuple[str, ...] = field(default_factory=tuple)
     assertion_ctes: tuple[CompileSqlTestCte, ...] = field(default_factory=tuple)

src/sqlbuild/compiler/planner/_helpers/resolve/refs.py

@@ -253,6 +253,42 @@ def resolve_table_function_references(
     return "".join(parts)
 
 
+def resolve_table_function_fixture_references(
+    *, query_sql: str, fixtures: dict[str, str]
+) -> tuple[str, frozenset[str]]:
+    """Replace complete table-function invocations with model-test fixture relations."""
+
+    parts: list[str] = []
+    last_index: int = 0
+    reached: set[str] = set()
+    match: re.Match[str]
+    for match in _iter_executable_matches(sql=query_sql, pattern=_TABLE_FUNCTION_PATTERN):
+        if match.start() < last_index:
+            continue
+        function_name: str = match.group(1)
+        fixture_sql: str | None = fixtures.get(function_name)
+        if fixture_sql is None:
+            continue
+        call_suffix_start: int = _skip_whitespace(sql=query_sql, start=match.end())
+        if (
+            call_suffix_start >= len(query_sql)
+            or query_sql[call_suffix_start] != SQL_FUNCTION_CALL_OPEN_PAREN
+        ):
+            continue
+        call_suffix_end: int = find_matching_paren(
+            sql=query_sql,
+            open_paren_index=call_suffix_start,
+            context="SQL table function fixture call",
+        )
+        parts.append(query_sql[last_index : match.start()])
+        parts.append(fixture_sql)
+        last_index = call_suffix_end + 1
+        reached.add(function_name)
+
+    parts.append(query_sql[last_index:])
+    return "".join(parts), frozenset(reached)
+
+
 def _iter_executable_matches(*, sql: str, pattern: re.Pattern[str]) -> Iterator[re.Match[str]]:
     index: int = 0
     while index < len(sql):