diff --git a/python/nwis_client/src/hydrotools/nwis_client/iv.py b/python/nwis_client/src/hydrotools/nwis_client/iv.py index b57fb544..fdc01b4d 100644 --- a/python/nwis_client/src/hydrotools/nwis_client/iv.py +++ b/python/nwis_client/src/hydrotools/nwis_client/iv.py @@ -118,6 +118,12 @@ def __init__(self, *, ) self._value_time_label = value_time_label + def __enter__(self): + return self + + def __exit__(self, *args): + self._restclient.close() + @verify_case_insensitive_kwargs(handler=_verify_case_insensitive_kwargs_handler) def get( self, diff --git a/python/nwis_client/tests/test_nwis.py b/python/nwis_client/tests/test_nwis.py index 120cc8b6..7f80d956 100644 --- a/python/nwis_client/tests/test_nwis.py +++ b/python/nwis_client/tests/test_nwis.py @@ -87,7 +87,7 @@ def wrap_iv_cache_location_to_temp(loop): o = partial(iv.IVDataService, cache_filename=cache_file) return o - + @pytest.fixture def setup_iv(IVDataServiceWithTempCache): @@ -119,7 +119,7 @@ def mocked_iv(mock_iv, setup_iv): `iv.IVDataService`'s `get_raw` method has been mocked to return an empty list. """ return setup_iv - + simplify_variable_test_data = [ ("test", ",", "test"), @@ -556,6 +556,17 @@ def test_nwis_client_cache_path(loop): # close resources service._restclient.close() + +@pytest.mark.slow +def test_nwis_client_context_manager(loop): + """verify that context manager closes resources""" + with iv.IVDataService() as service: + service.get(sites=["01189000"], startDT="2022-01-01") + assert service._restclient._session.closed + # should this be closed? The assertion fails if uncommented + # assert service._restclient._loop.is_closed() + + def test_fixes_209(loop, monkeypatch): """ verify that pandas FutureWarning is not raised by `IVDataService.get`. This FutureWarning was