Test socket types

Author: technige

neo4j/v1/connection.py

@@ -324,11 +324,7 @@ def connect(host, port=None, **config):
     s = create_connection((host, port))
 
     # Secure the connection if so requested
-    try:
-        secure = environ["NEO4J_SECURE"]
-    except KeyError:
-        secure = config.get("secure", False)
-    if secure:
+    if config.get("secure", False):
         if __debug__: log_info("~~ [SECURE] %s", host)
         s = secure_socket(s, host)
 

test/test_session.py

@@ -19,6 +19,8 @@
 # limitations under the License.
 
 
+from socket import socket
+from ssl import SSLSocket
 from unittest import TestCase
 
 from mock import patch
@@ -60,9 +62,6 @@ def test_session_that_dies_in_the_pool_will_not_be_given_out(self):
         session_2 = driver.session()
         assert session_2 is not session_1
 
-
-class RunTestCase(TestCase):
-
     def test_must_use_valid_url_scheme(self):
         with self.assertRaises(ValueError):
             GraphDatabase.driver("x://xxx")
@@ -83,6 +82,21 @@ def test_sessions_are_not_reused_if_still_in_use(self):
         session_1.close()
         assert session_1 is not session_2
 
+    def test_insecure_session_uses_insecure_socket(self):
+        driver = GraphDatabase.driver("bolt://localhost", secure=False)
+        session = driver.session()
+        assert isinstance(session.connection.channel.socket, socket)
+        session.close()
+
+    def test_secure_session_uses_secure_socket(self):
+        driver = GraphDatabase.driver("bolt://localhost", secure=True)
+        session = driver.session()
+        assert isinstance(session.connection.channel.socket, SSLSocket)
+        session.close()
+
+
+class RunTestCase(TestCase):
+
     def test_can_run_simple_statement(self):
         session = GraphDatabase.driver("bolt://localhost").session()
         count = 0