diff --git a/osmapi/OsmApi.py b/osmapi/OsmApi.py index 3c903b7..ec636cd 100644 --- a/osmapi/OsmApi.py +++ b/osmapi/OsmApi.py @@ -251,13 +251,26 @@ def __init__( self._session = self._get_http_session() def __del__(self): + self.close() + + return None + + def __enter__(self): + self._session = self._get_http_session() + return self + + def __exit__(self, *args): + self.close() + + def close(self): try: if self._changesetauto: self._changesetautoflush(True) except ResponseEmptyApiError: pass - return None + if self._session: + self._session.close() ################################################## # Capabilities # diff --git a/tests/helper_tests.py b/tests/helper_tests.py index a8ec2a3..defebad 100644 --- a/tests/helper_tests.py +++ b/tests/helper_tests.py @@ -23,6 +23,7 @@ def setupMock(self, status=200): mock_response.reason = "test reason" mock_response.content = 'test response' self.api._session.request = mock.Mock(return_value=mock_response) + self.api._session.close = mock.Mock() self.api._username = 'testuser' self.api._password = 'testpassword' @@ -56,6 +57,15 @@ def test_passwordfile_with_colon(self): self.assertEquals('testuser', my_api._username) self.assertEquals('test:userpass', my_api._password) + def test_close_call(self): + self.api.close() + self.assertEquals(self.api._session.close.call_count, 1) + + def test_close_context_manager(self): + with osmapi.OsmApi() as my_api: + my_api._session.close = mock.Mock() + self.assertEquals(my_api._session.close.call_count, 1) + def test_http_request_get(self): response = self.api._http_request( 'GET',