use patch context manager

Author: sfc-gh-mkubik

test/unit/test_session_manager.py

@@ -45,71 +45,71 @@ def create_session(
 
 @mock.patch("snowflake.connector.network.SnowflakeRestful.make_requests_session")
 def test_no_url_multiple_sessions(make_session_mock):
-    mock_conn = mock.Mock(name="test_no_url_multiple_sessions::mock_conn")
-    mock_conn.disable_request_pooling = False
-    mock_conn._ocsp_mode = lambda: DEFAULT_OCSP_MODE
+    with mock.patch("snowflake.connector.SnowflakeConnection") as mock_conn:
+        mock_conn.disable_request_pooling = False
+        mock_conn._ocsp_mode = lambda: DEFAULT_OCSP_MODE
 
-    rest = SnowflakeRestful(connection=mock_conn)
+        rest = SnowflakeRestful(connection=mock_conn)
 
-    create_session(rest, 2)
+        create_session(rest, 2)
 
-    assert make_session_mock.call_count == 2
+        assert make_session_mock.call_count == 2
 
-    assert list(rest._sessions_map.keys()) == [None]
+        assert list(rest._sessions_map.keys()) == [None]
 
-    session_pool = rest._sessions_map[None]
-    assert len(session_pool._idle_sessions) == 2
-    assert len(session_pool._active_sessions) == 0
+        session_pool = rest._sessions_map[None]
+        assert len(session_pool._idle_sessions) == 2
+        assert len(session_pool._active_sessions) == 0
 
-    close_sessions(rest, 1)
+        close_sessions(rest, 1)
 
 
 @mock.patch("snowflake.connector.network.SnowflakeRestful.make_requests_session")
 def test_multiple_urls_multiple_sessions(make_session_mock):
-    mock_conn = mock.Mock(name="test_multiple_urls_multiple_sessions::mock_conn")
-    mock_conn.disable_request_pooling = False
-    mock_conn._ocsp_mode = lambda: DEFAULT_OCSP_MODE
+    with mock.patch("snowflake.connector.SnowflakeConnection") as mock_conn:
+        mock_conn.disable_request_pooling = False
+        mock_conn._ocsp_mode = lambda: DEFAULT_OCSP_MODE
 
-    rest = SnowflakeRestful(connection=mock_conn)
+        rest = SnowflakeRestful(connection=mock_conn)
 
-    for url in [url_1, url_2, None]:
-        create_session(rest, num_sessions=2, url=url)
+        for url in [url_1, url_2, None]:
+            create_session(rest, num_sessions=2, url=url)
 
-    assert make_session_mock.call_count == 6
+        assert make_session_mock.call_count == 6
 
-    hostnames = list(rest._sessions_map.keys())
-    for hostname in [hostname_1, hostname_2, None]:
-        assert hostname in hostnames
+        hostnames = list(rest._sessions_map.keys())
+        for hostname in [hostname_1, hostname_2, None]:
+            assert hostname in hostnames
 
-    for pool in rest._sessions_map.values():
-        assert len(pool._idle_sessions) == 2
-        assert len(pool._active_sessions) == 0
+        for pool in rest._sessions_map.values():
+            assert len(pool._idle_sessions) == 2
+            assert len(pool._active_sessions) == 0
 
-    close_sessions(rest, 3)
+        close_sessions(rest, 3)
 
 
 @mock.patch("snowflake.connector.network.SnowflakeRestful.make_requests_session")
 def test_multiple_urls_reuse_sessions(make_session_mock):
-    mock_conn = mock.Mock(name="test_multiple_urls_reuse_sessions::mock_conn")
-    mock_conn.disable_request_pooling = False
-    mock_conn._ocsp_mode = lambda: DEFAULT_OCSP_MODE
+    with mock.patch("snowflake.connector.SnowflakeConnection") as mock_conn:
+        mock_conn.disable_request_pooling = False
+        mock_conn._ocsp_mode = lambda: DEFAULT_OCSP_MODE
 
-    rest = SnowflakeRestful(connection=mock_conn)
-    for url in [url_1, url_2, url_3, None]:
-        # create 10 sessions, one after another
-        for _ in range(10):
-            create_session(rest, url=url)
+        rest = SnowflakeRestful(connection=mock_conn)
+        for url in [url_1, url_2, url_3, None]:
+            # create 10 sessions, one after another
+            for _ in range(10):
+                create_session(rest, url=url)
 
-    # only one session is created and reused thereafter
-    assert make_session_mock.call_count == 3
+        # only one session is created and reused thereafter
+        assert make_session_mock.call_count == 3
 
-    hostnames = list(rest._sessions_map.keys())
-    assert len(hostnames) == 3
-    for hostname in [hostname_1, hostname_2, None]:
-        assert hostname in hostnames
+        hostnames = list(rest._sessions_map.keys())
+        assert len(hostnames) == 3
+        for hostname in [hostname_1, hostname_2, None]:
+            assert hostname in hostnames
 
-    for pool in rest._sessions_map.values():
-        assert len(pool._idle_sessions) == 1
-        assert len(pool._active_sessions) == 0
+        for pool in rest._sessions_map.values():
+            assert len(pool._idle_sessions) == 1
+            assert len(pool._active_sessions) == 0
 
-    close_sessions(rest, 3)
+        close_sessions(rest, 3)