From d309ac40159d70a092e16d3149690dec8ad1c1c3 Mon Sep 17 00:00:00 2001 From: Tatiana Al-Chueyr Date: Mon, 26 Feb 2024 12:26:53 +0000 Subject: [PATCH] Extend DatabricksTokenProfileMapping test to include session properties An Astronomer customer reported an issue while setting session properties in the DatabricksTokenProfileMapping class. We are extending the tests to include an example of how this works at the moment. --- tests/profiles/databricks/test_dbr_token.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/profiles/databricks/test_dbr_token.py b/tests/profiles/databricks/test_dbr_token.py index 3fd36a784c..0a1701af25 100644 --- 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(