DATA-1339 Set application name based on streamlit being in sys.modules (snowflakedb#1117)

sfc-gh-zblackwood · sfc-gh-mkeller · web-flow · commit 7826d3ede756 · 2022-05-10T10:27:12.000-07:00
Co-authored-by: Mark Keller &lt;63477823+sfc-gh-mkeller@users.noreply.github.com&gt;
diff --git a/src/snowflake/connector/connection.py b/src/snowflake/connector/connection.py
@@ -268,8 +268,11 @@ def __init__(self, **kwargs):
 
         self.heartbeat_thread = None
 
-        if "application" not in kwargs and ENV_VAR_PARTNER in os.environ.keys():
-            kwargs["application"] = os.environ[ENV_VAR_PARTNER]
+        if "application" not in kwargs:
+            if ENV_VAR_PARTNER in os.environ.keys():
+                kwargs["application"] = os.environ[ENV_VAR_PARTNER]
+            elif "streamlit" in sys.modules:
+                kwargs["application"] = "streamlit"
 
         self.converter = None
         self.__set_error_attributes()
diff --git a/test/unit/test_connection.py b/test/unit/test_connection.py
@@ -6,6 +6,7 @@
 
 import json
 import os
+import sys
 from unittest.mock import patch
 
 import pytest
@@ -19,45 +20,43 @@
     QueryStatus = None
 
 
-@patch("snowflake.connector.network.SnowflakeRestful._post_request")
-def test_connect_with_service_name(mockSnowflakeRestfulPostRequest):
-    def mock_post_request(url, headers, json_body, **kwargs):
-        global mock_cnt
-        ret = None
-        if mock_cnt == 0:
-            # return from /v1/login-request
-            ret = {
-                "success": True,
-                "message": None,
-                "data": {
-                    "token": "TOKEN",
-                    "masterToken": "MASTER_TOKEN",
-                    "idToken": None,
-                    "parameters": [
-                        {"name": "SERVICE_NAME", "value": "FAKE_SERVICE_NAME"}
-                    ],
-                },
-            }
-        return ret
+def fake_connector() -> snowflake.connector.SnowflakeConnection:
+    return snowflake.connector.connect(
+        user="user",
+        account="account",
+        password="testpassword",
+        database="TESTDB",
+        warehouse="TESTWH",
+    )
 
-    # POST requests mock
-    mockSnowflakeRestfulPostRequest.side_effect = mock_post_request
 
-    global mock_cnt
-    mock_cnt = 0
+@pytest.fixture
+def mock_post_requests(monkeypatch):
+    request_body = {}
 
-    account = "testaccount"
-    user = "testuser"
+    def mock_post_request(request, url, headers, json_body, **kwargs):
+        nonlocal request_body
+        request_body.update(json.loads(json_body))
+        return {
+            "success": True,
+            "message": None,
+            "data": {
+                "token": "TOKEN",
+                "masterToken": "MASTER_TOKEN",
+                "idToken": None,
+                "parameters": [{"name": "SERVICE_NAME", "value": "FAKE_SERVICE_NAME"}],
+            },
+        }
 
-    # connection
-    con = snowflake.connector.connect(
-        account=account,
-        user=user,
-        password="testpassword",
-        database="TESTDB",
-        warehouse="TESTWH",
+    monkeypatch.setattr(
+        snowflake.connector.network.SnowflakeRestful, "_post_request", mock_post_request
     )
-    assert con.service_name == "FAKE_SERVICE_NAME"
+
+    return request_body
+
+
+def test_connect_with_service_name(mock_post_requests):
+    assert fake_connector().service_name == "FAKE_SERVICE_NAME"
 
 
 @pytest.mark.skip(reason="Mock doesn't work as expected.")
@@ -137,38 +136,22 @@ def test_is_still_running():
 
 
 @pytest.mark.skipolddriver
-@patch("snowflake.connector.network.SnowflakeRestful._post_request")
-def test_partner_env_var(mockSnowflakeRestfulPostRequest):
+def test_partner_env_var(mock_post_requests):
     PARTNER_NAME = "Amanda"
 
-    request_body = {}
+    with patch.dict(os.environ, {ENV_VAR_PARTNER: PARTNER_NAME}):
+        assert fake_connector().application == PARTNER_NAME
 
-    def mock_post_request(url, headers, json_body, **kwargs):
-        nonlocal request_body
-        request_body = json.loads(json_body)
-        return {
-            "success": True,
-            "message": None,
-            "data": {
-                "token": "TOKEN",
-                "masterToken": "MASTER_TOKEN",
-                "idToken": None,
-                "parameters": [{"name": "SERVICE_NAME", "value": "FAKE_SERVICE_NAME"}],
-            },
-        }
+    assert (
+        mock_post_requests["data"]["CLIENT_ENVIRONMENT"]["APPLICATION"] == PARTNER_NAME
+    )
 
-    # POST requests mock
-    mockSnowflakeRestfulPostRequest.side_effect = mock_post_request
 
-    with patch.dict(os.environ, {ENV_VAR_PARTNER: PARTNER_NAME}):
-        # connection
-        con = snowflake.connector.connect(
-            user="user",
-            account="account",
-            password="testpassword",
-            database="TESTDB",
-            warehouse="TESTWH",
-        )
-        assert con.application == PARTNER_NAME
+@pytest.mark.skipolddriver
+def test_imported_module(mock_post_requests):
+    with patch.dict(sys.modules, {"streamlit": "foo"}):
+        assert fake_connector().application == "streamlit"
 
-    assert request_body["data"]["CLIENT_ENVIRONMENT"]["APPLICATION"] == PARTNER_NAME
+    assert (
+        mock_post_requests["data"]["CLIENT_ENVIRONMENT"]["APPLICATION"] == "streamlit"
+    )
