Extend DatabricksTokenProfileMapping test to include session proper…

…ties (#858) An Astronomer customer reported an issue setting session properties in the `DatabricksTokenProfileMapping` class. We are extending the tests to include an example of the current behaviour.
astronomer · Feb 27, 2024 · 5d0ad74 · 5d0ad74
1 parent aca9a6b
commit 5d0ad74
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/tests/profiles/databricks/test_dbr_token.py b/tests/profiles/databricks/test_dbr_token.py
@@ -95,11 +95,15 @@ def test_profile_args(
         profile_args={
             "schema": "my_schema",
             "catalog": "my_catalog",
+            "session_properties": {"legacy_time_parser_policy": "corrected"},
+            "threads": 4,
         },
     )
     assert profile_mapping.profile_args == {
         "schema": "my_schema",
         "catalog": "my_catalog",
+        "session_properties": {"legacy_time_parser_policy": "corrected"},
+        "threads": 4,
     }
 
     assert profile_mapping.profile == {
@@ -109,7 +113,23 @@ def test_profile_args(
         "http_path": mock_databricks_conn.extra_dejson.get("http_path"),
         "schema": "my_schema",
         "catalog": "my_catalog",
+        "threads": 4,
+        "session_properties": {"legacy_time_parser_policy": "corrected"},
     }
+    expected_profile_yml = """example:
+    outputs:
+        cosmos_target:
+            catalog: my_catalog
+            host: my_host
+            http_path: my_http_path
+            schema: my_schema
+            session_properties:
+                legacy_time_parser_policy: corrected
+            threads: 4
+            token: '{{ env_var(''COSMOS_CONN_DATABRICKS_TOKEN'') }}'
+            type: databricks
+    target: cosmos_target\n"""
+    assert profile_mapping.get_profile_file_contents("example") == expected_profile_yml
 
 
 def test_profile_args_overrides(