Update tests to check for batch_context

dbt-labs · Nov 27, 2024 · 345e474 · 345e474
1 parent c925282
commit 345e474
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 1 deletion.
diff --git a/tests/functional/microbatch/test_microbatch.py b/tests/functional/microbatch/test_microbatch.py
@@ -498,6 +498,13 @@ def test_run_with_event_time(self, project):
 {{ config(materialized='incremental', incremental_strategy='microbatch', unique_key='id', event_time='event_time', batch_size='day', begin=modules.datetime.datetime(2020, 1, 1, 0, 0, 0)) }}
 {{ log("start: "~ model.config.__dbt_internal_microbatch_event_time_start, info=True)}}
 {{ log("end: "~ model.config.__dbt_internal_microbatch_event_time_end, info=True)}}
+{% if model.batch_context %}
+{{ log("batch_context.event_time_start: "~ model.batch_context.event_time_start, info=True)}}
+{{ log("batch_context.event_time_end: "~ model.batch_context.event_time_end, info=True)}}
+{{ log("batch_context.id: "~ model.batch_context.id, info=True)}}
+{{ log("start timezone: "~ model.batch_context.event_time_start.tzinfo, info=True)}}
+{{ log("end timezone: "~ model.batch_context.event_time_end.tzinfo, info=True)}}
+{% endif %}
 select * from {{ ref('input_model') }}
 """
 
@@ -516,12 +523,23 @@ def test_run_with_event_time_logs(self, project):
 
         assert "start: 2020-01-01 00:00:00+00:00" in logs
         assert "end: 2020-01-02 00:00:00+00:00" in logs
+        assert "batch_context.event_time_start: 2020-01-01 00:00:00+00:00" in logs
+        assert "batch_context.event_time_end: 2020-01-02 00:00:00+00:00" in logs
+        assert "batch_context.id: 20200101" in logs
+        assert "start timezone: UTC" in logs
+        assert "end timezone: UTC" in logs
 
         assert "start: 2020-01-02 00:00:00+00:00" in logs
         assert "end: 2020-01-03 00:00:00+00:00" in logs
+        assert "batch_context.event_time_start: 2020-01-02 00:00:00+00:00" in logs
+        assert "batch_context.event_time_end: 2020-01-03 00:00:00+00:00" in logs
+        assert "batch_context.id: 20200102" in logs
 
         assert "start: 2020-01-03 00:00:00+00:00" in logs
         assert "end: 2020-01-03 13:57:00+00:00" in logs
+        assert "batch_context.event_time_start: 2020-01-03 00:00:00+00:00" in logs
+        assert "batch_context.event_time_end: 2020-01-03 13:57:00+00:00" in logs
+        assert "batch_context.id: 20200103" in logs
 
 
 microbatch_model_failing_incremental_partition_sql = """

diff --git a/tests/unit/contracts/graph/test_manifest.py b/tests/unit/contracts/graph/test_manifest.py
@@ -96,6 +96,7 @@
         "deprecation_date",
         "defer_relation",
         "time_spine",
+        "batch_context",
     }
 )
 

diff --git a/tests/unit/materializations/incremental/test_microbatch.py b/tests/unit/materializations/incremental/test_microbatch.py
@@ -605,7 +605,6 @@ def test_truncate_timestamp(self, timestamp, batch_size, expected_timestamp):
     @pytest.mark.parametrize(
         "batch_size,batch_start,expected_formatted_batch_start",
         [
-            (None, None, None),
             (BatchSize.year, datetime(2020, 1, 1, 1), "2020-01-01"),
             (BatchSize.month, datetime(2020, 1, 1, 1), "2020-01-01"),
             (BatchSize.day, datetime(2020, 1, 1, 1), "2020-01-01"),