From 78be77f62d672ebf6b51818e58cc981871b96be8 Mon Sep 17 00:00:00 2001
From: Courtney Holcomb <courtneyeholcomb@gmail.com>
Date: Fri, 28 Jun 2024 08:15:54 -0700
Subject: [PATCH] Actually set default_granularity during transformation

---
 .../transformations/default_granularity.py    |  2 ++
 .../test_configurable_transform_rules.py      | 33 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/dbt_semantic_interfaces/transformations/default_granularity.py b/dbt_semantic_interfaces/transformations/default_granularity.py
index 6e5f9af6..0e6b75fd 100644
--- a/dbt_semantic_interfaces/transformations/default_granularity.py
+++ b/dbt_semantic_interfaces/transformations/default_granularity.py
@@ -45,4 +45,6 @@ def transform_model(semantic_manifest: PydanticSemanticManifest) -> PydanticSema
                     ):
                         default_granularity = dimension.type_params.time_granularity
 
+            metric.default_granularity = default_granularity
+
         return semantic_manifest
diff --git a/tests/transformations/test_configurable_transform_rules.py b/tests/transformations/test_configurable_transform_rules.py
index 573d0d71..bdd5fefc 100644
--- a/tests/transformations/test_configurable_transform_rules.py
+++ b/tests/transformations/test_configurable_transform_rules.py
@@ -1,12 +1,18 @@
+from typing import Dict
+
 from dbt_semantic_interfaces.implementations.semantic_manifest import (
     PydanticSemanticManifest,
 )
+from dbt_semantic_interfaces.transformations.default_granularity import (
+    SetDefaultGranularityRule,
+)
 from dbt_semantic_interfaces.transformations.semantic_manifest_transformer import (
     PydanticSemanticManifestTransformer,
 )
 from dbt_semantic_interfaces.transformations.transform_rule import (
     SemanticManifestTransformRule,
 )
+from dbt_semantic_interfaces.type_enums import TimeGranularity
 
 
 class SliceNamesRule(SemanticManifestTransformRule):
@@ -32,3 +38,30 @@ def test_can_configure_model_transform_rules(  # noqa: D
     rules = [SliceNamesRule()]
     transformed_model = PydanticSemanticManifestTransformer.transform(pre_model, ordered_rule_sequences=(rules,))
     assert all(len(x.name) == 3 for x in transformed_model.semantic_models)
+
+
+def test_set_default_granularity_rule(  # noqa: D
+    simple_semantic_manifest__with_primary_transforms: PydanticSemanticManifest,
+) -> None:
+    pre_model = simple_semantic_manifest__with_primary_transforms
+
+    metric_exists_without_default_granularity = False
+    configured_default_granularities: Dict[str, TimeGranularity] = {}
+    for metric in pre_model.metrics:
+        if metric.default_granularity:
+            configured_default_granularities[metric.name] = metric.default_granularity
+            metric_exists_without_default_granularity = True
+
+    assert (
+        pre_model.metrics and metric_exists_without_default_granularity
+    ), "If there are no metrics without a configured default_granularity, this tests nothing."
+
+    rules = [SetDefaultGranularityRule()]
+    transformed_model = PydanticSemanticManifestTransformer.transform(pre_model, ordered_rule_sequences=(rules,))
+
+    for metric in transformed_model.metrics:
+        assert metric.default_granularity, f"No default_granularity set in transformation for metric '{metric.name}'"
+        if metric.name in configured_default_granularities:
+            assert (
+                metric.default_granularity == configured_default_granularities[metric.name]
+            ), f"Default granularity was unexpected changed during transformation for metric '{metric.name}"