Fix unit tests

snowflakedb · Apr 8, 2024 · 1f8b3f4 · 1f8b3f4
1 parent 864d07b
commit 1f8b3f4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/src/snowflake/cli/app/snow_connector.py b/src/snowflake/cli/app/snow_connector.py
@@ -78,7 +78,10 @@ def connect_to_snowflake(
             ] = diag_allowlist_path
 
     if (
-        overrides["session_token"] and overrides["master_token"]
+        "session_token" in overrides
+        and overrides["session_token"] is not None
+        and "master_token" in overrides
+        and overrides["master_token"] is not None
     ):  # TODO (lmonteromarin): Confirm this works as intended
         connection_parameters["server_session_keep_alive"] = True
 

diff --git a/tests/test_connection.py b/tests/test_connection.py
@@ -474,6 +474,7 @@ def test_session_and_master_tokens(mock_connector, mock_ctx, runner):
         database="test_dv",
         schema="PUBLIC",
         warehouse="xsmall",
+        server_session_keep_alive=True,
     )