Migrate data set conversion test to use snapshots.

dbt-labs · Sep 19, 2023 · e2b55bb · e2b55bb
1 parent 17fcee0
commit e2b55bb
Showing 1 changed file with 11 additions and 135 deletions.
diff --git a/metricflow/test/dataset/test_convert_semantic_model.py b/metricflow/test/dataset/test_convert_semantic_model.py
@@ -4,20 +4,12 @@
 
 from _pytest.fixtures import FixtureRequest
 from dbt_semantic_interfaces.references import SemanticModelReference
-from dbt_semantic_interfaces.type_enums.time_granularity import TimeGranularity
 
 from metricflow.protocols.sql_client import SqlClient
-from metricflow.specs.specs import (
-    DimensionSpec,
-    EntityReference,
-    EntitySpec,
-    InstanceSpecSet,
-    MeasureSpec,
-    TimeDimensionSpec,
-)
 from metricflow.sql.render.sql_plan_renderer import SqlQueryPlanRenderer
 from metricflow.test.fixtures.model_fixtures import ConsistentIdObjectRepository
 from metricflow.test.fixtures.setup_fixtures import MetricFlowTestSessionState
+from metricflow.test.snapshot_utils import assert_spec_set_snapshot_equal
 from metricflow.test.sql.compare_sql_plan import assert_rendered_sql_equal
 
 logger = logging.getLogger(__name__)
@@ -32,52 +24,12 @@ def test_convert_table_semantic_model_without_measures(  # noqa: D
     """Simple test for converting a table semantic model. Since there are no measures, primary time is not checked."""
     users_data_set = consistent_id_object_repository.simple_model_data_sets["users_latest"]
 
-    expected_spec_set = InstanceSpecSet(
-        metric_specs=(),
-        measure_specs=(),
-        dimension_specs=(
-            DimensionSpec(element_name="home_state_latest", entity_links=()),
-            DimensionSpec(
-                element_name="home_state_latest",
-                entity_links=(EntityReference(element_name="user"),),
-            ),
-        ),
-        entity_specs=(EntitySpec(element_name="user", entity_links=()),),
-        time_dimension_specs=(
-            TimeDimensionSpec(element_name="ds", entity_links=(), time_granularity=TimeGranularity.DAY),
-            TimeDimensionSpec(element_name="ds", entity_links=(), time_granularity=TimeGranularity.WEEK),
-            TimeDimensionSpec(element_name="ds", entity_links=(), time_granularity=TimeGranularity.MONTH),
-            TimeDimensionSpec(element_name="ds", entity_links=(), time_granularity=TimeGranularity.QUARTER),
-            TimeDimensionSpec(element_name="ds", entity_links=(), time_granularity=TimeGranularity.YEAR),
-            TimeDimensionSpec(
-                element_name="ds",
-                entity_links=(EntityReference(element_name="user"),),
-                time_granularity=TimeGranularity.DAY,
-            ),
-            TimeDimensionSpec(
-                element_name="ds",
-                entity_links=(EntityReference(element_name="user"),),
-                time_granularity=TimeGranularity.WEEK,
-            ),
-            TimeDimensionSpec(
-                element_name="ds",
-                entity_links=(EntityReference(element_name="user"),),
-                time_granularity=TimeGranularity.MONTH,
-            ),
-            TimeDimensionSpec(
-                element_name="ds",
-                entity_links=(EntityReference(element_name="user"),),
-                time_granularity=TimeGranularity.QUARTER,
-            ),
-            TimeDimensionSpec(
-                element_name="ds",
-                entity_links=(EntityReference(element_name="user"),),
-                time_granularity=TimeGranularity.YEAR,
-            ),
-        ),
+    assert_spec_set_snapshot_equal(
+        request=request,
+        mf_test_session_state=mf_test_session_state,
+        set_id="result0",
+        spec_set=users_data_set.instance_set.spec_set,
     )
-
-    assert users_data_set.instance_set.spec_set == expected_spec_set
     assert users_data_set.semantic_model_reference == SemanticModelReference(semantic_model_name="users_latest")
     assert_rendered_sql_equal(
         request=request,
@@ -101,89 +53,13 @@ def test_convert_table_semantic_model_with_measures(  # noqa: D
     """
     id_verifications_data_set = consistent_id_object_repository.simple_model_data_sets["id_verifications"]
 
-    expected_spec_set = InstanceSpecSet(
-        metric_specs=(),
-        measure_specs=(MeasureSpec(element_name="identity_verifications"),),
-        dimension_specs=(
-            DimensionSpec(element_name="verification_type", entity_links=()),
-            DimensionSpec(
-                element_name="verification_type",
-                entity_links=(EntityReference(element_name="verification"),),
-            ),
-        ),
-        entity_specs=(
-            EntitySpec(element_name="verification", entity_links=()),
-            EntitySpec(element_name="user", entity_links=()),
-            EntitySpec(
-                element_name="user",
-                entity_links=(EntityReference(element_name="verification"),),
-            ),
-        ),
-        time_dimension_specs=(
-            TimeDimensionSpec(element_name="ds", entity_links=(), time_granularity=TimeGranularity.DAY),
-            TimeDimensionSpec(element_name="ds", entity_links=(), time_granularity=TimeGranularity.WEEK),
-            TimeDimensionSpec(element_name="ds", entity_links=(), time_granularity=TimeGranularity.MONTH),
-            TimeDimensionSpec(element_name="ds", entity_links=(), time_granularity=TimeGranularity.QUARTER),
-            TimeDimensionSpec(element_name="ds", entity_links=(), time_granularity=TimeGranularity.YEAR),
-            TimeDimensionSpec(element_name="ds_partitioned", entity_links=(), time_granularity=TimeGranularity.DAY),
-            TimeDimensionSpec(element_name="ds_partitioned", entity_links=(), time_granularity=TimeGranularity.WEEK),
-            TimeDimensionSpec(element_name="ds_partitioned", entity_links=(), time_granularity=TimeGranularity.MONTH),
-            TimeDimensionSpec(element_name="ds_partitioned", entity_links=(), time_granularity=TimeGranularity.QUARTER),
-            TimeDimensionSpec(element_name="ds_partitioned", entity_links=(), time_granularity=TimeGranularity.YEAR),
-            TimeDimensionSpec(
-                element_name="ds",
-                entity_links=(EntityReference(element_name="verification"),),
-                time_granularity=TimeGranularity.DAY,
-            ),
-            TimeDimensionSpec(
-                element_name="ds",
-                entity_links=(EntityReference(element_name="verification"),),
-                time_granularity=TimeGranularity.WEEK,
-            ),
-            TimeDimensionSpec(
-                element_name="ds",
-                entity_links=(EntityReference(element_name="verification"),),
-                time_granularity=TimeGranularity.MONTH,
-            ),
-            TimeDimensionSpec(
-                element_name="ds",
-                entity_links=(EntityReference(element_name="verification"),),
-                time_granularity=TimeGranularity.QUARTER,
-            ),
-            TimeDimensionSpec(
-                element_name="ds",
-                entity_links=(EntityReference(element_name="verification"),),
-                time_granularity=TimeGranularity.YEAR,
-            ),
-            TimeDimensionSpec(
-                element_name="ds_partitioned",
-                entity_links=(EntityReference(element_name="verification"),),
-                time_granularity=TimeGranularity.DAY,
-            ),
-            TimeDimensionSpec(
-                element_name="ds_partitioned",
-                entity_links=(EntityReference(element_name="verification"),),
-                time_granularity=TimeGranularity.WEEK,
-            ),
-            TimeDimensionSpec(
-                element_name="ds_partitioned",
-                entity_links=(EntityReference(element_name="verification"),),
-                time_granularity=TimeGranularity.MONTH,
-            ),
-            TimeDimensionSpec(
-                element_name="ds_partitioned",
-                entity_links=(EntityReference(element_name="verification"),),
-                time_granularity=TimeGranularity.QUARTER,
-            ),
-            TimeDimensionSpec(
-                element_name="ds_partitioned",
-                entity_links=(EntityReference(element_name="verification"),),
-                time_granularity=TimeGranularity.YEAR,
-            ),
-        ),
+    assert_spec_set_snapshot_equal(
+        request=request,
+        mf_test_session_state=mf_test_session_state,
+        set_id="result0",
+        spec_set=id_verifications_data_set.instance_set.spec_set,
     )
 
-    assert id_verifications_data_set.instance_set.spec_set == expected_spec_set
     assert id_verifications_data_set.semantic_model_reference == SemanticModelReference(
         semantic_model_name="id_verifications"
     )