Skip to content

Commit 05a4acd

Browse files
Add explicit close() for session. Add context manager support.
Fixes #278
1 parent 1c260f2 commit 05a4acd

2 files changed

Lines changed: 20 additions & 0 deletions

File tree

prometheus_api_client/prometheus_connect.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,17 @@ def __init__(
9494
self._session.proxies = proxy
9595
self._session.mount(self.url, HTTPAdapter(max_retries=retry))
9696

97+
def close(self) -> None:
98+
"""Close the underlying HTTP session and release connection-pool resources."""
99+
self._session.close()
100+
101+
def __enter__(self):
102+
"""Support usage as a context manager: ``with PrometheusConnect(...) as pc:``"""
103+
return self
104+
105+
def __exit__(self, exc_type, exc_val, exc_tb):
106+
self.close()
107+
97108
def check_prometheus_connection(self, params: dict = None) -> bool:
98109
"""
99110
Check Promethus connection.

tests/test_prometheus_connect.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -395,5 +395,14 @@ def test_get_label_values_method(self): # noqa D102
395395
request = handler.requests[0]
396396
self.assertEqual(request.path_url, "/api/v1/label/label_name/values")
397397

398+
def test_close(self): # noqa D102
399+
self.pc.close() # must not raise
400+
401+
def test_context_manager(self): # noqa D102
402+
with PrometheusConnect(url="http://doesnt_matter.xyz", disable_ssl=True) as pc:
403+
self.assertIsInstance(pc, PrometheusConnect)
404+
pc.close() # idempotent — must not raise after __exit__
405+
406+
398407
if __name__ == "__main__":
399408
unittest.main()

0 commit comments

Comments
 (0)