Use MocketSocket as context manager

mocket/__init__.py

@@ -2,4 +2,4 @@
 
 __all__ = ("mocketize", "Mocket", "MocketEntry", "Mocketizer")
 
-__version__ = "3.9.38"
+__version__ = "3.9.39"

mocket/mocket.py

@@ -133,7 +133,6 @@ def __init__(
         self, family=socket.AF_INET, type=socket.SOCK_STREAM, proto=0, *args, **kwargs
     ):
         self.true_socket = true_socket(family, type, proto)
-        self._connected = False
         self._buflen = 65536
         self._entry = None
         self.family = int(family)
@@ -160,6 +159,12 @@ def __str__(self):  # pragma: no cover
             self.__class__.__name__, self.family, self.type, self.proto
         )
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
     @property
     def fd(self):
         if self._fd is None:
@@ -378,6 +383,8 @@ def send(self, data, *args, **kwargs):  # pragma: no cover
         return len(data)
 
     def close(self):
+        if self.true_socket and not self.true_socket._closed:
+            self.true_socket.close()
         self._fd = None
 
     def __getattr__(self, name):

tests/main/test_mocket.py

@@ -127,6 +127,16 @@ def test_makefile(self):
             self.assertEqual(fp.read().strip(), encode_to_bytes("Show me."))
             self.assertEqual(len(Mocket._requests), 1)
 
+    def test_socket_as_context_manager(self):
+        addr = ("localhost", 80)
+        Mocket.register(MocketEntry(addr, ["Show me.\r\n"]))
+        with Mocketizer():
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as _so:
+                _so.connect(addr)
+                _so.sendall(encode_to_bytes("Whatever..."))
+                data = _so.recv(4096)
+                self.assertEqual(data, encode_to_bytes("Show me.\r\n"))
+
 
 class MocketizeTestCase(TestCase):
     def mocketize_setup(self):