Create metricflow-semantics Package

metricflow-semantics/metricflow_semantics/test_helpers/snapshot_helpers.py

@@ -6,9 +6,19 @@
 import re
 import webbrowser
 from dataclasses import dataclass
-from typing import Callable, List, Optional, Tuple
+from typing import Callable, List, Optional, Tuple, TypeVar
 
 import _pytest.fixtures
+import tabulate
+from _pytest.fixtures import FixtureRequest
+
+from metricflow_semantics.dag.mf_dag import MetricFlowDag
+from metricflow_semantics.mf_logging.pretty_print import mf_pformat
+from metricflow_semantics.model.semantics.linkable_element_set import LinkableElementSet
+from metricflow_semantics.naming.object_builder_scheme import ObjectBuilderNamingScheme
+from metricflow_semantics.specs.spec_classes import InstanceSpecSet, LinkableSpecSet
+from metricflow_semantics.test_helpers.config_helpers import MetricFlowTestConfiguration
+from tests_metricflow.snapshot_utils import assert_object_snapshot_equal, assert_str_snapshot_equal
 
 logger = logging.getLogger(__name__)
 
@@ -165,3 +175,167 @@ def add_overwrite_snapshots_cli_flag(parser: _pytest.config.argparsing.Parser) -
         action="store_true",
         help="Overwrites existing snapshots by ones generated during this testing session",
     )
+
+
+# In plan outputs, replace strings that vary from run to run with this so that comparisons can be made
+# consistently.
+PLACEHOLDER_CHAR_FOR_INCOMPARABLE_STRINGS = "*"
+
+
+def make_schema_replacement_function(system_schema: str, source_schema: str) -> Callable[[str], str]:
+    """Generates a function to replace schema names in test outputs."""
+
+    # The schema of the warehouse used in tests changes from run to run, so don't compare those.
+    def replacement_function(text: str) -> str:
+        # Replace with a string of the same length so that indents are preserved.
+        text = text.replace(source_schema, PLACEHOLDER_CHAR_FOR_INCOMPARABLE_STRINGS * len(source_schema))
+        # Same with the MetricFlow system schema.
+        return text.replace(system_schema, PLACEHOLDER_CHAR_FOR_INCOMPARABLE_STRINGS * len(system_schema))
+
+    return replacement_function
+
+
+def replace_dataset_id_hash(text: str) -> str:
+    """Replaces data set ID hashes for primed semantic models.
+
+    The data set ID hash changes from run to run because it's based on the DW schema in the semantic model, which changes
+    run to run.
+    """
+    pattern = re.compile(r"'[a-zA-Z0-9_]+__[a-zA-Z0-9_]+__(?P<hash>[a-zA-Z0-9_]+)'")
+    while True:
+        match = pattern.search(text)
+        if match:
+            data_set_id_hash = match.group("hash")
+            # Replace with the same length to preserve indents
+            text = text.replace(data_set_id_hash, PLACEHOLDER_CHAR_FOR_INCOMPARABLE_STRINGS * len(data_set_id_hash))
+        else:
+            break
+    return text
+
+
+PlanT = TypeVar("PlanT", bound=MetricFlowDag)
+
+
+def assert_plan_snapshot_text_equal(
+    request: FixtureRequest,
+    mf_test_configuration: MetricFlowTestConfiguration,
+    plan: PlanT,
+    plan_snapshot_text: str,
+    plan_snapshot_file_extension: str = ".xml",
+    exclude_line_regex: Optional[str] = None,
+    incomparable_strings_replacement_function: Optional[Callable[[str], str]] = None,
+    additional_sub_directories_for_snapshots: Tuple[str, ...] = (),
+) -> None:
+    """Checks if the given plan text is equal to the one that's saved for comparison.
+
+    * The location of the file is automatically generated based on the test and the plan's ID.
+    * This may create a new saved plan file or overwrite the existing one, depending on the configuration.
+    * replace_incomparable_strings is used to replace strings in the plan text before comparison. Useful for making
+      plans consistent when there are strings that vary between runs and shouldn't be compared.
+    * additional_sub_directories_for_snapshots is used to specify additional sub-directories (in the automatically
+      generated directory) where plan outputs should reside.
+
+    TODO: Make this more generic by renaming plan -> DAG.
+    """
+    assert_snapshot_text_equal(
+        request=request,
+        mf_test_configuration=mf_test_configuration,
+        group_id=plan.__class__.__name__,
+        snapshot_id=str(plan.dag_id),
+        snapshot_text=plan_snapshot_text,
+        snapshot_file_extension=plan_snapshot_file_extension,
+        exclude_line_regex=exclude_line_regex,
+        incomparable_strings_replacement_function=incomparable_strings_replacement_function,
+        additional_sub_directories_for_snapshots=additional_sub_directories_for_snapshots,
+    )
+
+
+def assert_linkable_element_set_snapshot_equal(  # noqa: D103
+    request: FixtureRequest,
+    mf_test_configuration: MetricFlowTestConfiguration,
+    set_id: str,
+    linkable_element_set: LinkableElementSet,
+) -> None:
+    headers = ("Semantic Model", "Entity Links", "Name", "Time Granularity", "Date Part", "Properties")
+    rows = []
+    for linkable_dimension_iterable in linkable_element_set.path_key_to_linkable_dimensions.values():
+        for linkable_dimension in linkable_dimension_iterable:
+            rows.append(
+                (
+                    # Checking a limited set of fields as the result is large due to the paths in the object.
+                    (
+                        linkable_dimension.semantic_model_origin.semantic_model_name
+                        if linkable_dimension.semantic_model_origin
+                        else None
+                    ),
+                    tuple(entity_link.element_name for entity_link in linkable_dimension.entity_links),
+                    linkable_dimension.element_name,
+                    linkable_dimension.time_granularity.name if linkable_dimension.time_granularity is not None else "",
+                    linkable_dimension.date_part.name if linkable_dimension.date_part is not None else "",
+                    sorted(
+                        linkable_element_property.name for linkable_element_property in linkable_dimension.properties
+                    ),
+                )
+            )
+
+    for linkable_entity_iterable in linkable_element_set.path_key_to_linkable_entities.values():
+        for linkable_entity in linkable_entity_iterable:
+            rows.append(
+                (
+                    # Checking a limited set of fields as the result is large due to the paths in the object.
+                    linkable_entity.semantic_model_origin.semantic_model_name,
+                    tuple(entity_link.element_name for entity_link in linkable_entity.entity_links),
+                    linkable_entity.element_name,
+                    "",
+                    "",
+                    sorted(linkable_element_property.name for linkable_element_property in linkable_entity.properties),
+                )
+            )
+
+    for linkable_metric_iterable in linkable_element_set.path_key_to_linkable_metrics.values():
+        for linkable_metric in linkable_metric_iterable:
+            rows.append(
+                (
+                    # Checking a limited set of fields as the result is large due to the paths in the object.
+                    linkable_metric.join_by_semantic_model.semantic_model_name,
+                    tuple(entity_link.element_name for entity_link in linkable_entity.entity_links),
+                    linkable_metric.element_name,
+                    "",
+                    "",
+                    sorted(linkable_element_property.name for linkable_element_property in linkable_metric.properties),
+                )
+            )
+
+    assert_str_snapshot_equal(
+        request=request,
+        mf_test_configuration=mf_test_configuration,
+        snapshot_id=set_id,
+        snapshot_str=tabulate.tabulate(headers=headers, tabular_data=sorted(rows)),
+    )
+
+
+def assert_spec_set_snapshot_equal(  # noqa: D103
+    request: FixtureRequest, mf_test_configuration: MetricFlowTestConfiguration, set_id: str, spec_set: InstanceSpecSet
+) -> None:
+    assert_object_snapshot_equal(
+        request=request,
+        mf_test_configuration=mf_test_configuration,
+        obj_id=set_id,
+        obj=sorted(spec.qualified_name for spec in spec_set.all_specs),
+    )
+
+
+def assert_linkable_spec_set_snapshot_equal(  # noqa: D103
+    request: FixtureRequest, mf_test_configuration: MetricFlowTestConfiguration, set_id: str, spec_set: LinkableSpecSet
+) -> None:
+    # TODO: This will be used in a later PR and this message will be removed.
+    naming_scheme = ObjectBuilderNamingScheme()
+    assert_snapshot_text_equal(
+        request=request,
+        mf_test_configuration=mf_test_configuration,
+        group_id=spec_set.__class__.__name__,
+        snapshot_id=set_id,
+        snapshot_text=mf_pformat(sorted(naming_scheme.input_str(spec) for spec in spec_set.as_tuple)),
+        snapshot_file_extension=".txt",
+        additional_sub_directories_for_snapshots=(),
+    )

tests_metricflow/dataflow/builder/test_cyclic_join.py

@@ -12,12 +12,12 @@
     MetricSpec,
 )
 from metricflow_semantics.test_helpers.config_helpers import MetricFlowTestConfiguration
+from metricflow_semantics.test_helpers.snapshot_helpers import assert_plan_snapshot_text_equal
 
 from metricflow.dataflow.builder.dataflow_plan_builder import DataflowPlanBuilder
 from tests_metricflow.dataflow_plan_to_svg import display_graph_if_requested
 from tests_metricflow.fixtures.manifest_fixtures import MetricFlowEngineTestFixture, SemanticManifestSetup
 from tests_metricflow.fixtures.sql_client_fixtures import sql_client  # noqa: F401, F403
-from tests_metricflow.snapshot_utils import assert_plan_snapshot_text_equal
 
 logger = logging.getLogger(__name__)
 

tests_metricflow/dataflow/builder/test_dataflow_plan_builder.py

@@ -21,11 +21,11 @@
     TimeDimensionSpec,
 )
 from metricflow_semantics.test_helpers.config_helpers import MetricFlowTestConfiguration
+from metricflow_semantics.test_helpers.snapshot_helpers import assert_plan_snapshot_text_equal
 
 from metricflow.dataflow.builder.dataflow_plan_builder import DataflowPlanBuilder
 from metricflow.dataset.dataset_classes import DataSet
 from tests_metricflow.dataflow_plan_to_svg import display_graph_if_requested
-from tests_metricflow.snapshot_utils import assert_plan_snapshot_text_equal
 from tests_metricflow.time.metric_time_dimension import MTD_SPEC_DAY, MTD_SPEC_MONTH, MTD_SPEC_QUARTER, MTD_SPEC_WEEK
 
 logger = logging.getLogger(__name__)

tests_metricflow/dataflow/builder/test_node_data_set.py

@@ -18,6 +18,7 @@
 )
 from metricflow_semantics.sql.sql_join_type import SqlJoinType
 from metricflow_semantics.test_helpers.config_helpers import MetricFlowTestConfiguration
+from metricflow_semantics.test_helpers.snapshot_helpers import assert_spec_set_snapshot_equal
 
 from metricflow.dataflow.builder.node_data_set import DataflowPlanNodeOutputDataSetResolver
 from metricflow.dataflow.nodes.join_to_base import JoinDescription, JoinToBaseOutputNode
@@ -33,7 +34,6 @@
 )
 from metricflow.sql.sql_table import SqlTable
 from tests_metricflow.fixtures.manifest_fixtures import MetricFlowEngineTestFixture, SemanticManifestSetup
-from tests_metricflow.snapshot_utils import assert_spec_set_snapshot_equal
 
 logger = logging.getLogger(__name__)
 

tests_metricflow/dataflow/optimizer/source_scan/test_cm_branch_combiner.py

@@ -8,6 +8,7 @@
 from metricflow_semantics.dag.mf_dag import DagId
 from metricflow_semantics.specs.spec_classes import InstanceSpecSet, MeasureSpec
 from metricflow_semantics.test_helpers.config_helpers import MetricFlowTestConfiguration
+from metricflow_semantics.test_helpers.snapshot_helpers import assert_plan_snapshot_text_equal
 
 from metricflow.dataflow.dataflow_plan import (
     BaseOutput,
@@ -21,7 +22,6 @@
 )
 from tests_metricflow.dataflow_plan_to_svg import display_graph_if_requested
 from tests_metricflow.fixtures.manifest_fixtures import MetricFlowEngineTestFixture, SemanticManifestSetup
-from tests_metricflow.snapshot_utils import assert_plan_snapshot_text_equal
 
 
 def make_dataflow_plan(node: BaseOutput) -> DataflowPlan:  # noqa: D103

tests_metricflow/dataflow/optimizer/source_scan/test_source_scan_optimizer.py

@@ -15,6 +15,7 @@
     MetricSpec,
 )
 from metricflow_semantics.test_helpers.config_helpers import MetricFlowTestConfiguration
+from metricflow_semantics.test_helpers.snapshot_helpers import assert_plan_snapshot_text_equal
 
 from metricflow.dataflow.builder.dataflow_plan_builder import DataflowPlanBuilder
 from metricflow.dataflow.dataflow_plan import (
@@ -43,7 +44,6 @@
 from metricflow.dataflow.optimizer.source_scan.source_scan_optimizer import SourceScanOptimizer
 from metricflow.dataset.dataset_classes import DataSet
 from tests_metricflow.dataflow_plan_to_svg import display_graph_if_requested
-from tests_metricflow.snapshot_utils import assert_plan_snapshot_text_equal
 
 logger = logging.getLogger(__name__)
 

tests_metricflow/dataset/test_convert_semantic_model.py

@@ -7,10 +7,10 @@
 from _pytest.fixtures import FixtureRequest
 from dbt_semantic_interfaces.references import SemanticModelReference
 from metricflow_semantics.test_helpers.config_helpers import MetricFlowTestConfiguration
+from metricflow_semantics.test_helpers.snapshot_helpers import assert_spec_set_snapshot_equal
 
 from metricflow.protocols.sql_client import SqlClient
 from tests_metricflow.fixtures.manifest_fixtures import MetricFlowEngineTestFixture, SemanticManifestSetup
-from tests_metricflow.snapshot_utils import assert_spec_set_snapshot_equal
 from tests_metricflow.sql.compare_sql_plan import assert_rendered_sql_equal
 
 logger = logging.getLogger(__name__)

tests_metricflow/plan_conversion/test_dataflow_to_sql_plan.py

@@ -30,6 +30,7 @@
 from metricflow_semantics.sql.sql_bind_parameters import SqlBindParameters
 from metricflow_semantics.sql.sql_join_type import SqlJoinType
 from metricflow_semantics.test_helpers.config_helpers import MetricFlowTestConfiguration
+from metricflow_semantics.test_helpers.snapshot_helpers import assert_plan_snapshot_text_equal
 
 from metricflow.dataflow.builder.dataflow_plan_builder import DataflowPlanBuilder
 from metricflow.dataflow.dataflow_plan import (
@@ -53,7 +54,6 @@
 from metricflow.sql.optimizer.optimization_levels import SqlQueryOptimizationLevel
 from tests_metricflow.dataflow_plan_to_svg import display_graph_if_requested
 from tests_metricflow.fixtures.manifest_fixtures import MetricFlowEngineTestFixture, SemanticManifestSetup
-from tests_metricflow.snapshot_utils import assert_plan_snapshot_text_equal
 from tests_metricflow.sql.compare_sql_plan import assert_rendered_sql_from_plan_equal, assert_sql_plan_text_equal
 from tests_metricflow.time.metric_time_dimension import MTD_SPEC_DAY
 

tests_metricflow/semantics/model/semantics/test_linkable_spec_resolver.py

@@ -21,8 +21,10 @@
 )
 from metricflow_semantics.model.semantics.semantic_model_join_evaluator import MAX_JOIN_HOPS
 from metricflow_semantics.test_helpers.config_helpers import MetricFlowTestConfiguration
-
-from tests_metricflow.snapshot_utils import assert_linkable_element_set_snapshot_equal, assert_spec_set_snapshot_equal
+from metricflow_semantics.test_helpers.snapshot_helpers import (
+    assert_linkable_element_set_snapshot_equal,
+    assert_spec_set_snapshot_equal,
+)
 
 logger = logging.getLogger(__name__)
 

tests_metricflow/semantics/model/test_semantic_model_container.py

@@ -10,8 +10,11 @@
 from metricflow_semantics.model.semantics.metric_lookup import MetricLookup
 from metricflow_semantics.model.semantics.semantic_model_lookup import SemanticModelLookup
 from metricflow_semantics.test_helpers.config_helpers import MetricFlowTestConfiguration
+from metricflow_semantics.test_helpers.snapshot_helpers import (
+    assert_linkable_element_set_snapshot_equal,
+)
 
-from tests_metricflow.snapshot_utils import assert_linkable_element_set_snapshot_equal, assert_object_snapshot_equal
+from tests_metricflow.snapshot_utils import assert_object_snapshot_equal
 
 logger = logging.getLogger(__name__)
 

tests_metricflow/semantics/query/group_by_item/resolution_dag/test_resolution_dags.py

@@ -8,9 +8,9 @@
 from metricflow_semantics.model.semantic_manifest_lookup import SemanticManifestLookup
 from metricflow_semantics.query.group_by_item.resolution_dag.dag import GroupByItemResolutionDag
 from metricflow_semantics.test_helpers.config_helpers import MetricFlowTestConfiguration
+from metricflow_semantics.test_helpers.snapshot_helpers import assert_plan_snapshot_text_equal
 
 from tests_metricflow.semantics.query.group_by_item.ambiguous_resolution_query_id import AmbiguousResolutionQueryId
-from tests_metricflow.snapshot_utils import assert_plan_snapshot_text_equal
 
 logger = logging.getLogger(__name__)
 

tests_metricflow/semantics/query/group_by_item/test_available_group_by_items.py

@@ -10,9 +10,9 @@
 from metricflow_semantics.query.group_by_item.resolution_dag.dag import GroupByItemResolutionDag
 from metricflow_semantics.specs.spec_classes import LinkableSpecSet
 from metricflow_semantics.test_helpers.config_helpers import MetricFlowTestConfiguration
+from metricflow_semantics.test_helpers.snapshot_helpers import assert_linkable_spec_set_snapshot_equal
 
 from tests_metricflow.semantics.query.group_by_item.conftest import AmbiguousResolutionQueryId
-from tests_metricflow.snapshot_utils import assert_linkable_spec_set_snapshot_equal
 
 logger = logging.getLogger(__name__)