feat(API): lock workspace in with block through lock() method

Author: yt-ms

CHANGELOG.rst

@@ -5,6 +5,7 @@ History
 Next Release
 ------------
 * Feat: Add support for grouping elements (#72)
+* Feat: Locking workspace in ``with`` block through ``lock()`` method (#62)
 
 0.4.0 (2021-02-05)
 ------------------

src/structurizr/api/structurizr_client.py

@@ -21,6 +21,7 @@
 import hmac
 import logging
 from base64 import b64encode
+from contextlib import contextmanager
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
 from typing import Dict
@@ -99,7 +100,11 @@ def __repr__(self) -> str:
         )
 
     def __enter__(self):
-        """Enter a context by locking the corresponding remote workspace."""
+        """Enter a context by locking the corresponding remote workspace.
+
+        Note: this method is deprecated in favour of lock(), and will be removed in a
+        future relesae.
+        """
         is_successful = self.lock_workspace()
         if not is_successful:
             raise StructurizrClientException(
@@ -109,14 +114,36 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        """Exit a context by unlocking the corresponding remote workspace."""
+        """Exit a context by unlocking the corresponding remote workspace.
+
+        Note: this method is deprecated in favour of lock(), and will be removed in a
+        future relesae.
+        """
         is_successful = self.unlock_workspace()
         self._client.close()
         if exc_type is None and not is_successful:
             raise StructurizrClientException(
                 f"Failed to unlock the Structurizr workspace {self.workspace_id}."
             )
 
+    @contextmanager
+    def lock(self):
+        """Provide a context manager for locking and unlocking a workspace."""
+        is_successful = self.lock_workspace()
+        if not is_successful:
+            raise StructurizrClientException(
+                f"Failed to lock the Structurizr workspace {self.workspace_id}."
+            )
+        try:
+            yield None
+        finally:
+            is_successful = self.unlock_workspace()
+            self._client.close()
+            if not is_successful:
+                raise StructurizrClientException(
+                    f"Failed to unlock the Structurizr workspace {self.workspace_id}."
+                )
+
     def close(self) -> None:
         """Close the connection pool."""
         self._client.close()

tests/unit/api/test_structurizr_client.py

@@ -210,3 +210,29 @@ def fake_send(request: Request):
     assert requests[0].url.path == "/workspace/19/lock"
     assert requests[1].method == "DELETE"
     assert requests[1].url.path == "/workspace/19/lock"
+
+
+def test_locking_and_unlocking_with_context_manager(
+    client: StructurizrClient, mocker: MockerFixture
+):
+    """Check new-style locking using .lock()."""
+    requests: List[Request] = []
+
+    def fake_send(request: Request):
+        nonlocal requests
+        requests.append(request)
+        return Response(
+            200,
+            content='{"success": true, "message": "OK"}'.encode("ascii"),
+            request=request,
+        )
+
+    mocker.patch.object(client._client, "send", new=fake_send)
+    with client.lock():
+        pass
+
+    assert len(requests) == 2
+    assert requests[0].method == "PUT"
+    assert requests[0].url.path == "/workspace/19/lock"
+    assert requests[1].method == "DELETE"
+    assert requests[1].url.path == "/workspace/19/lock"