refactor: MocketSSLSocket use proper ssl-context instead of urllib3

mindflayer · Nov 25, 2024 · e485148 · e485148
1 parent 1cce25a
commit e485148
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git a/mocket/ssl/socket.py b/mocket/ssl/socket.py
@@ -62,16 +62,23 @@ def unwrap(self) -> MocketSocket:
         return self._original_socket
 
     @classmethod
-    def _create(cls, sock: MocketSocket, *args: Any, **kwargs: Any) -> MocketSSLSocket:
+    def _create(
+        cls,
+        sock: MocketSocket,
+        ssl_context: ssl.SSLContext | None = None,
+        server_hostname: str | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> MocketSSLSocket:
         ssl_socket = MocketSSLSocket()
         ssl_socket._original_socket = sock
+        ssl_socket._true_socket = sock._true_socket
 
-        from mocket.urllib3 import true_ssl_wrap_socket
-
-        ssl_socket._true_socket = true_ssl_wrap_socket(
-            sock._true_socket,
-            **kwargs,
-        )
+        if ssl_context:
+            ssl_socket._true_socket = ssl_context.wrap_socket(
+                sock=ssl_socket._true_socket,
+                server_hostname=server_hostname,
+            )
 
         ssl_socket._kwargs = kwargs
 

diff --git a/mocket/urllib3.py b/mocket/urllib3.py
@@ -2,13 +2,10 @@
 
 from typing import Any
 
-import urllib3
 from mocket.socket import MocketSocket
 from mocket.ssl.context import MocketSSLContext
 from mocket.ssl.socket import MocketSSLSocket
 
-true_ssl_wrap_socket = urllib3.util.ssl_.ssl_wrap_socket
-
 
 def mock_match_hostname(*args: Any) -> None:
     return None