Make Client support context manager protocol

CHANGELOG.md

@@ -7,7 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-- n/a
+### Added
+
+- `Client` instances can now be used in a `with` statement to manage the lifecycle
+  of the underlying threads.
 
 ## [2.11.0] - 2021-07-15
 

README.md

@@ -29,15 +29,15 @@ Usage Example
 from pubtools.pulplib import Client
 
 # Make a client pointing at this Pulp server
-client = Client(url='https://pulp.example.com/', auth=('admin', 'some-password'))
+with Client(url='https://pulp.example.com/', auth=('admin', 'some-password')) as client:
 
-# Get a particular repo by ID.
-# All methods return Future instances; .result() blocks
-repo = client.get_repository('zoo').result()
+  # Get a particular repo by ID.
+  # All methods return Future instances; .result() blocks
+  repo = client.get_repository('zoo').result()
 
-# Pulp objects have relevant methods, e.g. publish().
-# Returned future may encapsulate one or more Pulp tasks.
-publish = repo.publish().result()
+  # Pulp objects have relevant methods, e.g. publish().
+  # Returned future may encapsulate one or more Pulp tasks.
+  publish = repo.publish().result()
 ```
 
 Development

docs/index.rst

@@ -28,12 +28,12 @@ the desired methods to perform actions on Pulp.
     from pubtools.pulplib import Client
 
     # Make a client pointing at this Pulp server
-    client = Client(url='https://pulp.example.com/', auth=('admin', 'some-password'))
+    with Client(url='https://pulp.example.com/', auth=('admin', 'some-password')) as client:
 
-    # Get a particular repo by ID.
-    # All methods return Future instances; .result() blocks
-    repo = client.get_repository('zoo').result()
+        # Get a particular repo by ID.
+        # All methods return Future instances; .result() blocks
+        repo = client.get_repository('zoo').result()
 
-    # Pulp objects have relevant methods, e.g. publish().
-    # Returned future may encapsulate one or more Pulp tasks.
-    publish = repo.publish().result()
+        # Pulp objects have relevant methods, e.g. publish().
+        # Returned future may encapsulate one or more Pulp tasks.
+        publish = repo.publish().result()

examples/upload-files

@@ -73,8 +73,8 @@ def main():
         logging.getLogger("pubtools.pulplib").setLevel(logging.DEBUG)
         log.setLevel(logging.DEBUG)
 
-    client = make_client(p)
-    return upload(client, p.path, p.repo_id)
+    with make_client(p) as client:
+        return upload(client, p.path, p.repo_id)
 
 
 if __name__ == "__main__":

pubtools/pulplib/_impl/client/client.py

@@ -41,6 +41,21 @@ class Client(object):
     If a future is currently awaiting one or more Pulp tasks, cancelling the future
     will attempt to cancel those tasks.
 
+    **Client lifecycle:**
+
+    .. versionadded:: 2.12.0
+
+    Client instances support the context manager protocol and can be used
+    via a ``with`` statement, as in example:
+
+    .. code-block:: python
+
+        with Client(url="https://pulp.example.com/") as client:
+            do_something_with(client)
+
+    While not mandatory, it is encouraged to ensure that any threads associated with
+    the client are promptly shut down.
+
     **Proxy futures:**
 
     .. versionadded:: 2.1.0
@@ -173,6 +188,13 @@ def __init__(self, url, **kwargs):
             .with_retry(retry_policy=self._RETRY_POLICY)
         )
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        self._request_executor.__exit__(*args, **kwargs)
+        self._task_executor.__exit__(*args, **kwargs)
+
     def get_repository(self, repository_id):
         """Get a repository by ID.
 

pubtools/pulplib/_impl/fake/client.py

@@ -69,8 +69,24 @@ def __init__(self):
         self._lock = threading.RLock()
         self._uuidgen = random.Random()
         self._uuidgen.seed(0)
+        self._shutdown = False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_args, **_kwargs):
+        self._shutdown = True
+
+    def _ensure_alive(self):
+        if self._shutdown:
+            # We are technically capable of working just fine after shutdown,
+            # but the point of this class is to be an accurate stand-in for
+            # a real client, so raise the same kind of exception here
+            raise RuntimeError("cannot schedule new futures after shutdown")
 
     def search_repository(self, criteria=None):
+        self._ensure_alive()
+
         criteria = criteria or Criteria.true()
         repos = []
 
@@ -94,6 +110,8 @@ def search_repository(self, criteria=None):
         return self._prepare_pages(repos)
 
     def search_content(self, criteria=None):
+        self._ensure_alive()
+
         criteria = criteria or Criteria.true()
         out = []
 
@@ -129,6 +147,8 @@ def search_content(self, criteria=None):
         return self._prepare_pages(out)
 
     def search_distributor(self, criteria=None):
+        self._ensure_alive()
+
         criteria = criteria or Criteria.true()
         distributors = []
 
@@ -205,13 +225,17 @@ def get_repository(self, repository_id):
         return f_proxy(f_return(data[0]))
 
     def get_maintenance_report(self):
+        self._ensure_alive()
+
         if self._maintenance_report:
             report = MaintenanceReport._from_data(json.loads(self._maintenance_report))
         else:
             report = MaintenanceReport()
         return f_proxy(f_return(report))
 
     def set_maintenance(self, report):
+        self._ensure_alive()
+
         report_json = json.dumps(report._export_dict(), indent=4, sort_keys=True)
         report_fileobj = StringIO(report_json)
 
@@ -226,6 +250,8 @@ def set_maintenance(self, report):
         return f_proxy(publish_ft)
 
     def get_content_type_ids(self):
+        self._ensure_alive()
+
         return f_proxy(f_return(self._type_ids))
 
     def _do_upload_file(self, upload_id, file_obj, name):

setup.py

@@ -21,7 +21,7 @@ def get_requirements():
 
 setup(
     name="pubtools-pulplib",
-    version="2.11.0",
+    version="2.12.0",
     packages=find_packages(exclude=["tests"]),
     package_data={"pubtools.pulplib._impl.schema": ["*.yaml"]},
     url="https://github.com/release-engineering/pubtools-pulplib",

tests/client/test_client.py

@@ -68,6 +68,26 @@ def test_can_search(client, requests_mocker):
     assert requests_mocker.call_count == 1
 
 
+def test_client_lifecycle(client, requests_mocker):
+    """Client is usable in with statement"""
+
+    requests_mocker.post(
+        "https://pulp.example.com/pulp/api/v2/repositories/search/",
+        json=[{"id": "repo1"}],
+    )
+
+    client = Client("https://pulp.example.com")
+    with client:
+        # This should work OK
+        assert client.search_repository().result()
+
+    # But after end of 'with' statement, it should be shut down
+    with pytest.raises(RuntimeError) as excinfo:
+        client.search_repository()
+
+    assert "cannot schedule new futures after shutdown" in str(excinfo.value)
+
+
 def test_can_search_distributor(client, requests_mocker):
     """search_distributor issues distributors/search POST request as expected."""
     requests_mocker.post(

tests/fake/test_fake.py

@@ -1,3 +1,5 @@
+from functools import partial
+
 import pytest
 
 from pubtools.pulplib import FakeController, Repository, PulpException
@@ -43,3 +45,27 @@ def test_get_wrong_type_raises():
     client = controller.client
     with pytest.raises(TypeError):
         client.get_repository(["oops", "should have been a string"])
+
+
+def test_client_lifecycle():
+    """FakeClient can be used in a with statement, and not afterwards."""
+    controller = FakeController()
+
+    with controller.client as client:
+        # This should work OK
+        assert client.search_repository().result()
+
+    # But after end of 'with' statement, most public methods will no longer work
+    for fn in [
+        client.search_repository,
+        client.search_content,
+        client.search_distributor,
+        partial(client.get_repository, "somerepo"),
+        client.get_maintenance_report,
+        partial(client.set_maintenance, {"what": "ever"}),
+        client.get_content_type_ids,
+    ]:
+        with pytest.raises(RuntimeError) as excinfo:
+            fn()
+
+        assert "cannot schedule new futures after shutdown" in str(excinfo.value)