Added context manager to Server class

oittaa · Nov 2, 2022 · 7dc4379 · 7dc4379
1 parent d2efa16
commit 7dc4379
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/src/gcp_storage_emulator/server.py b/src/gcp_storage_emulator/server.py
@@ -407,6 +407,14 @@ def __init__(self, host, port, in_memory, default_bucket=None, data_dir=None):
             buckets.create_bucket(default_bucket, self._storage)
         self._api = APIThread(host, port, self._storage)
 
+    # Context Manager
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, *args):
+        self.stop()
+
     def start(self):
         self._api.start()
         self._api.is_running.wait()  # Start the API thread