diff --git a/google_auth_oauthlib/flow.py b/google_auth_oauthlib/flow.py index c5d8bce..e564ca4 100644 --- a/google_auth_oauthlib/flow.py +++ b/google_auth_oauthlib/flow.py @@ -433,31 +433,33 @@ def run_local_server( bind_addr or host, port, wsgi_app, handler_class=_WSGIRequestHandler ) - redirect_uri_format = ( - "http://{}:{}/" if redirect_uri_trailing_slash else "http://{}:{}" - ) - self.redirect_uri = redirect_uri_format.format(host, local_server.server_port) - auth_url, _ = self.authorization_url(**kwargs) - - if open_browser: - # if browser is None it defaults to default browser - webbrowser.get(browser).open(auth_url, new=1, autoraise=True) - - if authorization_prompt_message: - print(authorization_prompt_message.format(url=auth_url)) - - local_server.timeout = timeout_seconds - local_server.handle_request() - - # Note: using https here because oauthlib is very picky that - # OAuth 2.0 should only occur over https. - authorization_response = wsgi_app.last_request_uri.replace("http", "https") - self.fetch_token( - authorization_response=authorization_response, audience=token_audience - ) - - # This closes the socket - local_server.server_close() + try: + redirect_uri_format = ( + "http://{}:{}/" if redirect_uri_trailing_slash else "http://{}:{}" + ) + self.redirect_uri = redirect_uri_format.format( + host, local_server.server_port + ) + auth_url, _ = self.authorization_url(**kwargs) + + if open_browser: + # if browser is None it defaults to default browser + webbrowser.get(browser).open(auth_url, new=1, autoraise=True) + + if authorization_prompt_message: + print(authorization_prompt_message.format(url=auth_url)) + + local_server.timeout = timeout_seconds + local_server.handle_request() + + # Note: using https here because oauthlib is very picky that + # OAuth 2.0 should only occur over https. + authorization_response = wsgi_app.last_request_uri.replace("http", "https") + self.fetch_token( + authorization_response=authorization_response, audience=token_audience + ) + finally: + local_server.server_close() return self.credentials diff --git a/tests/unit/test_flow.py b/tests/unit/test_flow.py index 3e61fcd..a3314d1 100644 --- a/tests/unit/test_flow.py +++ b/tests/unit/test_flow.py @@ -25,6 +25,7 @@ import pytest import requests import urllib +import webbrowser from google_auth_oauthlib import flow @@ -440,3 +441,17 @@ def test_run_local_server_occupied_port( with pytest.raises(OSError) as exc: instance.run_local_server(port=port) assert "address already in use" in exc.strerror.lower() + + @mock.patch("google_auth_oauthlib.flow.webbrowser.get", autospec=True) + @mock.patch("wsgiref.simple_server.make_server", autospec=True) + def test_local_server_socket_cleanup( + self, make_server_mock, webbrowser_mock, instance + ): + server_mock = mock.MagicMock() + make_server_mock.return_value = server_mock + webbrowser_mock.side_effect = webbrowser.Error("Browser not found") + + with pytest.raises(webbrowser.Error): + instance.run_local_server() + + server_mock.server_close.assert_called_once()