From 9f29106737972bfb2f4aea42b47938810cd5b985 Mon Sep 17 00:00:00 2001
From: Will Deng <william.deng14@gmail.com>
Date: Tue, 5 Nov 2024 13:01:49 -0500
Subject: [PATCH] added tests for parsing custom grains in where filters

---
 tests/example_project_configuration.py        |  8 +++++
 .../project_configuration.yaml                |  8 +++++
 tests/parsing/test_saved_query_parsing.py     | 29 +++++++++++++++++++
 tests/parsing/test_where_filter_parsing.py    | 16 ++++++++++
 .../test_where_filters_are_parseable.py       |  2 ++
 5 files changed, 63 insertions(+)

diff --git a/tests/example_project_configuration.py b/tests/example_project_configuration.py
index 71f407bf..c107165e 100644
--- a/tests/example_project_configuration.py
+++ b/tests/example_project_configuration.py
@@ -51,6 +51,14 @@
               primary_column:
                 name: ds_day
                 time_granularity: day
+            - node_relation:
+                schema_name: stuffs
+                alias: week_time_spine
+              primary_column:
+                name: ds
+                time_granularity: week
+              custom_granularities:
+                - name: martian_week
         """
     ),
 )
diff --git a/tests/fixtures/semantic_manifest_yamls/simple_semantic_manifest/project_configuration.yaml b/tests/fixtures/semantic_manifest_yamls/simple_semantic_manifest/project_configuration.yaml
index 80c6f34a..efce1c6a 100644
--- a/tests/fixtures/semantic_manifest_yamls/simple_semantic_manifest/project_configuration.yaml
+++ b/tests/fixtures/semantic_manifest_yamls/simple_semantic_manifest/project_configuration.yaml
@@ -11,3 +11,11 @@ project_configuration:
       primary_column:
         name: ds_day
         time_granularity: day
+    - node_relation:
+        alias: mf_time_spine
+        schema_name: stufffs
+      primary_column:
+        name: ds
+        time_granularity: day
+      custom_granularities:
+        - name: martian_day
diff --git a/tests/parsing/test_saved_query_parsing.py b/tests/parsing/test_saved_query_parsing.py
index 20b8f4b3..2bc04d11 100644
--- a/tests/parsing/test_saved_query_parsing.py
+++ b/tests/parsing/test_saved_query_parsing.py
@@ -120,6 +120,35 @@ def test_saved_query_group_by() -> None:
     )
 
 
+def test_saved_query_group_by_with_custom_grain() -> None:
+    """Test for parsing group_bys in a saved query."""
+    yaml_contents = textwrap.dedent(
+        """\
+        saved_query:
+          name: test_saved_query_group_bys
+          query_params:
+              metrics:
+                - test_metric_a
+              group_by:
+                - TimeDimension('test_entity__metric_time', 'martian_week')
+                - Dimension('test_entity__metric_time__martian_week')
+
+        """
+    )
+    file = YamlConfigFile(filepath="test_dir/inline_for_test", contents=yaml_contents)
+
+    build_result = parse_yaml_files_to_semantic_manifest(files=[file, EXAMPLE_PROJECT_CONFIGURATION_YAML_CONFIG_FILE])
+
+    assert len(build_result.semantic_manifest.saved_queries) == 1
+    saved_query = build_result.semantic_manifest.saved_queries[0]
+    assert len(saved_query.query_params.group_by) == 2
+    print(saved_query.query_params.group_by)
+    assert {
+        "TimeDimension('test_entity__metric_time', 'martian_week')",
+        "Dimension('test_entity__metric_time__martian_week')",
+    } == set(saved_query.query_params.group_by)
+
+
 def test_saved_query_where() -> None:
     """Test for parsing where clause in a saved query."""
     where = "Dimension(test_entity__test_dimension) == true"
diff --git a/tests/parsing/test_where_filter_parsing.py b/tests/parsing/test_where_filter_parsing.py
index 60a37e0a..0886ea72 100644
--- a/tests/parsing/test_where_filter_parsing.py
+++ b/tests/parsing/test_where_filter_parsing.py
@@ -196,6 +196,22 @@ def test_dimension_date_part() -> None:  # noqa
                 time_granularity_name=TimeGranularity.WEEK.value,
             ),
         ),
+        (
+            "{{ TimeDimension('metric_time__martian_week') }} > '2023-01-01'",
+            TimeDimensionCallParameterSet(
+                time_dimension_reference=TimeDimensionReference("metric_time"),
+                entity_path=(),
+                time_granularity_name="martian_week",
+            ),
+        ),
+        (
+            "{{ TimeDimension('metric_time', time_granularity_name='martian_week') }} > '2023-01-01'",
+            TimeDimensionCallParameterSet(
+                time_dimension_reference=TimeDimensionReference("metric_time"),
+                entity_path=(),
+                time_granularity_name="martian_week",
+            ),
+        ),
     ],
 )
 def test_time_dimension_grain(  # noqa
diff --git a/tests/validations/test_where_filters_are_parseable.py b/tests/validations/test_where_filters_are_parseable.py
index ae948a99..3bedd8ee 100644
--- a/tests/validations/test_where_filters_are_parseable.py
+++ b/tests/validations/test_where_filters_are_parseable.py
@@ -166,6 +166,7 @@ def test_metric_where_filter_validations_invalid_granularity(  # noqa: D
             PydanticWhereFilter(where_sql_template="{{ TimeDimension('metric_time', 'cool') }}"),
             PydanticWhereFilter(where_sql_template="{{ TimeDimension('metric_time', 'month') }}"),
             PydanticWhereFilter(where_sql_template="{{ TimeDimension('metric_time', 'MONTH') }}"),
+            PydanticWhereFilter(where_sql_template="{{ TimeDimension('metric_time', 'martian_day') }}"),
         ]
     )
     validator = SemanticManifestValidator[PydanticSemanticManifest]([WhereFiltersAreParseable()])
@@ -195,6 +196,7 @@ def test_saved_query_with_happy_filter(  # noqa: D
                 where=PydanticWhereFilterIntersection(
                     where_filters=[
                         PydanticWhereFilter(where_sql_template="{{ TimeDimension('metric_time', 'hour') }}"),
+                        PydanticWhereFilter(where_sql_template="{{ TimeDimension('metric_time', 'martian_day') }}"),
                     ]
                 ),
             ),