From bbe34cab20a4d1fcdf0138cb9f52dbfd07cc02f3 Mon Sep 17 00:00:00 2001 From: Jakob Schlyter Date: Wed, 20 Nov 2024 09:12:49 +0100 Subject: [PATCH] add support for contextlib --- dnstapir/key_resolver.py | 15 ++++++++++++++- tests/test_key_resolver.py | 15 ++++++++++++++- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/dnstapir/key_resolver.py b/dnstapir/key_resolver.py index b3ceb5e..807c214 100644 --- a/dnstapir/key_resolver.py +++ b/dnstapir/key_resolver.py @@ -94,6 +94,19 @@ def httpx_client(self) -> httpx.Client: self._httpx_client = httpx.Client() return self._httpx_client + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + def __del__(self): + self.close() + + def close(self): + """Explicitly close the client and free resources.""" if self._httpx_client is not None: - self._httpx_client.close() + try: + self._httpx_client.close() + finally: + self._httpx_client = None diff --git a/tests/test_key_resolver.py b/tests/test_key_resolver.py index 656c7c6..48528d9 100644 --- a/tests/test_key_resolver.py +++ b/tests/test_key_resolver.py @@ -15,6 +15,19 @@ def test_url_key_resolver(httpx_mock: HTTPXMock): httpx_mock.add_response(url=f"https://keys/{key_id}.pem", content=public_key_pem) resolver = UrlKeyResolver(client_database_base_url="https://keys") - res = resolver.resolve_public_key(key_id) assert res == public_key + + +def test_url_key_resolver_contextlib(httpx_mock: HTTPXMock): + key_id = "xyzzy" + public_key = ed25519.Ed25519PrivateKey.generate().public_key() + public_key_pem = public_key.public_bytes( + encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo + ) + + httpx_mock.add_response(url=f"https://keys/{key_id}.pem", content=public_key_pem) + + with UrlKeyResolver(client_database_base_url="https://keys") as resolver: + res = resolver.resolve_public_key(key_id) + assert res == public_key