From 3eb2c95f3f3c4ba7f5bf37b35eec494cabb914b8 Mon Sep 17 00:00:00 2001
From: Paul Yang <github@ardorful.com>
Date: Wed, 11 Dec 2024 12:59:24 -0800
Subject: [PATCH] /* PR_START p--misc 01 */ Add marker for DuckDb-only tests.

---
 tests_metricflow/fixtures/setup_fixtures.py | 27 ++++++++++++++++++---
 tests_metricflow/test_pytest_markers.py     | 23 ++++++++++++++++++
 2 files changed, 46 insertions(+), 4 deletions(-)
 create mode 100644 tests_metricflow/test_pytest_markers.py

diff --git a/tests_metricflow/fixtures/setup_fixtures.py b/tests_metricflow/fixtures/setup_fixtures.py
index e3835840f2..4b68d1188b 100644
--- a/tests_metricflow/fixtures/setup_fixtures.py
+++ b/tests_metricflow/fixtures/setup_fixtures.py
@@ -18,6 +18,7 @@
     add_overwrite_snapshots_cli_flag,
 )
 
+from metricflow.protocols.sql_client import SqlClient, SqlEngine
 from tests_metricflow import TESTS_METRICFLOW_DIRECTORY_ANCHOR
 from tests_metricflow.snapshots import METRICFLOW_SNAPSHOT_DIRECTORY_ANCHOR
 from tests_metricflow.table_snapshot.table_snapshots import SqlTableSnapshotHash, SqlTableSnapshotRepository
@@ -56,13 +57,20 @@ def pytest_addoption(parser: _pytest.config.argparsing.Parser) -> None:
 
 # Name of the pytest marker for tests that generate SQL-engine specific snapshots.
 SQL_ENGINE_SNAPSHOT_MARKER_NAME = "sql_engine_snapshot"
+# Name of the pytest marker to indicate that the test should only be run when the test session is configured to use
+# DuckDB as the SQL engine.
+DUCKDB_ONLY_MARKER_NAME = "duckdb_only"
 
 
 def pytest_configure(config: _pytest.config.Config) -> None:
     """Hook as specified by the pytest API for configuration."""
     config.addinivalue_line(
         name="markers",
-        line=f"{SQL_ENGINE_SNAPSHOT_MARKER_NAME}: mark tests as generating a snapshot specific to a SQL engine",
+        line=f"{SQL_ENGINE_SNAPSHOT_MARKER_NAME}: mark tests as generating a snapshot specific to a SQL engine.",
+    )
+    config.addinivalue_line(
+        name="markers",
+        line=f"{DUCKDB_ONLY_MARKER_NAME}: mark tests as one that should only be run with DuckDB.",
     )
     config.addinivalue_line(
         name="markers",
@@ -72,14 +80,25 @@ def pytest_configure(config: _pytest.config.Config) -> None:
 
 def check_sql_engine_snapshot_marker(request: FixtureRequest) -> None:
     """Raises an error if the given test request does not have the sql-engine-test marker set."""
-    mark_names = set(mark.name for mark in request.node.iter_markers(name=SQL_ENGINE_SNAPSHOT_MARKER_NAME))
-    if SQL_ENGINE_SNAPSHOT_MARKER_NAME not in mark_names:
+    if request.node.get_closest_marker(SQL_ENGINE_SNAPSHOT_MARKER_NAME) is None:
         raise ValueError(
-            f"This test needs to be marked with '{SQL_ENGINE_SNAPSHOT_MARKER_NAME}' to keep track of all tests that "
+            f"This test needs to be marked with {SQL_ENGINE_SNAPSHOT_MARKER_NAME!r} to keep track of all tests that "
             f"generate SQL-engine specific snapshots."
         )
 
 
+@pytest.fixture(autouse=True)
+def skip_if_not_duck_db(request: FixtureRequest, sql_client: SqlClient) -> None:
+    """Skip tests if the test is marked for testing with DuckDB only, but the test session uses a different engine."""
+    if request.node.get_closest_marker(DUCKDB_ONLY_MARKER_NAME) is None:
+        return
+
+    test_session_engine_type = sql_client.sql_engine_type
+    duckdb_engine_type = SqlEngine.DUCKDB
+    if test_session_engine_type is not duckdb_engine_type:
+        pytest.skip(f"Skipping as this test is only run with {duckdb_engine_type}, but {test_session_engine_type=}")
+
+
 @pytest.fixture(scope="session")
 def mf_test_configuration(  # noqa: D103
     request: FixtureRequest,
diff --git a/tests_metricflow/test_pytest_markers.py b/tests_metricflow/test_pytest_markers.py
new file mode 100644
index 0000000000..39dfd596bf
--- /dev/null
+++ b/tests_metricflow/test_pytest_markers.py
@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import logging
+
+import pytest
+
+from metricflow.protocols.sql_client import SqlClient, SqlEngine
+
+logger = logging.getLogger(__name__)
+
+
+@pytest.mark.duckdb_only
+def test_duckdb_only(sql_client: SqlClient) -> None:
+    """Check that the `duckdb_only` skips tests when the tests are run with another SQL engine.
+
+    This depends on auto-using of the `skip_if_not_duck_db` fixture.
+    """
+    duckdb_engine_type = SqlEngine.DUCKDB
+    engine_type_under_test = sql_client.sql_engine_type
+    assert engine_type_under_test is duckdb_engine_type, (
+        f"This test should have only been run when using {duckdb_engine_type}, but the engine type under test is"
+        f" {engine_type_under_test}"
+    )