Merge pull request #41 from didx-xyz/feature/context_manager

Feature/context manager

Author: morrieinmaas

aries_cloudcontroller/aries_controller_base.py

@@ -1,7 +1,7 @@
 from aiohttp import (
     ClientSession,
 )
-from abc import ABC
+from contextlib import AbstractAsyncContextManager
 from dataclasses import dataclass
 from pubsub import pub
 
@@ -27,7 +27,7 @@
 
 
 @dataclass
-class AriesAgentControllerBase(ABC):
+class AriesAgentControllerBase(AbstractAsyncContextManager):
     """The Aries Agent Controller class
 
     This class allows you to interact with Aries by exposing the aca-py API.
@@ -97,6 +97,12 @@ def __post_init__(self):
 
         self.revocations = RevocationController(self.admin_url, self.client_session)
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, exc_traceback):
+        await self.terminate()
+
     def init_webhook_server(self):
         raise NotImplementedError
 

setup.py

@@ -22,7 +22,7 @@ def parse_requirements(filename):
 if __name__ == "__main__":
     setup(
         name=PACKAGE_NAME,
-        version="0.2.6",
+        version="0.2.7",
         description="A simple python package for controlling an aries agent through the admin-api interface",
         long_description=long_description,
         long_description_content_type="text/markdown",

tests/test_aries_controller_base.py

@@ -211,3 +211,19 @@ async def test_terminate(self, caplog):
         with pytest.raises(AttributeError):
             assert ac.webhook_server is None
         await ac.terminate()
+
+    @pytest.mark.asyncio
+    async def test_context_manager(self, caplog):
+        caplog.set_level(logging.INFO)
+        api_key = "123456789"
+
+        async with AriesAgentControllerBase(admin_url="", api_key=api_key) as ac:
+            assert ac.client_session
+            assert ac.api_key == api_key
+            assert ac.headers == {"X-API-Key": api_key}
+            ac.remove_api_key()
+
+        assert ac.headers == {}
+        assert ac.api_key is None
+        assert ac.client_session.headers == {}
+        assert "Client Session closed." in caplog.text