From 2b62f4b7407aa7b1e6c5be4a214528f4262f47a8 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Wed, 15 May 2019 16:32:43 +0300 Subject: [PATCH 01/28] Initial implementation for headless login --- README.md | 21 ++++++++ neuromation/api/__init__.py | 9 ++++ neuromation/api/config_factory.py | 27 ++++++++++ neuromation/api/login.py | 83 ++++++++++++++++++++++++++----- neuromation/cli/config.py | 27 ++++++++++ 5 files changed, 154 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index a3ac33ddc..e2bd6207f 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ * [neuro config](#neuro-config) * [neuro config login](#neuro-config-login) * [neuro config login-with-token](#neuro-config-login-with-token) + * [neuro config login-headless](#neuro-config-login-headless) * [neuro config show](#neuro-config-show) * [neuro config show-token](#neuro-config-show-token) * [neuro config docker](#neuro-config-docker) @@ -722,6 +723,7 @@ Name | Description| |---|---| | _[neuro config login](#neuro-config-login)_| Log into Neuromation Platform | | _[neuro config login\-with-token](#neuro-config-login-with-token)_| Log into Neuromation Platform with token | +| _[neuro config login-headless](#neuro-config-login-headless)_| Log into Neuromation Platform from non-GUI server environment | | _[neuro config show](#neuro-config-show)_| Print current settings | | _[neuro config show-token](#neuro-config-show-token)_| Print current authorization token | | _[neuro config docker](#neuro-config-docker)_| Configure docker client for working with platform registry | @@ -768,6 +770,25 @@ Name | Description| +### neuro config login-headless + +Log into Neuromation Platform from non-GUI server environment.

URL is a platform entrypoint URL.

The command works similar to "neuro login" but instead of opening a browser
for performing OAuth registration prints an URL that should be open on guest
host.

Then user inputs a code displayed in a browser after successful login back
in neuro command to finish the login process. + +**Usage:** + +```bash +neuro config login-headless [OPTIONS] [URL] +``` + +**Options:** + +Name | Description| +|----|------------| +|_--help_|Show this message and exit.| + + + + ### neuro config show Print current settings. diff --git a/neuromation/api/__init__.py b/neuromation/api/__init__.py index c20385b2a..422ec04f4 100644 --- a/neuromation/api/__init__.py +++ b/neuromation/api/__init__.py @@ -159,5 +159,14 @@ async def login_with_token( await Factory(path).login_with_token(token, url=url, timeout=timeout) +async def login_headless( + *, + url: URL = DEFAULT_API_URL, + path: Optional[Path] = None, + timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT +) -> None: + await Factory(path).login_headless(url=url, timeout=timeout) + + async def logout(*, path: Optional[Path] = None) -> None: await Factory(path).logout() diff --git a/neuromation/api/config_factory.py b/neuromation/api/config_factory.py index 42580bccc..f94984b8d 100644 --- a/neuromation/api/config_factory.py +++ b/neuromation/api/config_factory.py @@ -13,6 +13,7 @@ from .core import DEFAULT_TIMEOUT from .login import ( AuthNegotiator, + HeadlessNegotiator, _AuthConfig, _AuthToken, _ClusterConfig, @@ -70,6 +71,30 @@ async def login( await client.jobs.list() # raises an exception if cannot login self._save(config) + async def login_headless( + self, + *, + url: URL = DEFAULT_API_URL, + timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT, + ) -> None: + if self._path.exists(): + raise ConfigError(f"Config file {self._path} already exists. Please logout") + config_unauthorized = await get_server_config(url) + negotiator = HeadlessNegotiator(config_unauthorized.auth_config) + auth_token = await negotiator.refresh_token() + + config_authorized = await get_server_config(url, token=auth_token.token) + config = _Config( + auth_config=config_authorized.auth_config, + auth_token=auth_token, + cluster_config=config_authorized.cluster_config, + pypi=_PyPIVersion.create_uninitialized(), + url=url, + ) + async with Client._create(config, timeout=timeout) as client: + await client.jobs.list() # raises an exception if cannot login + self._save(config) + async def login_with_token( self, token: str, @@ -139,6 +164,7 @@ def _serialize_auth_config(self, auth_config: _AuthConfig) -> Dict[str, Any]: "token_url": str(auth_config.token_url), "client_id": auth_config.client_id, "audience": auth_config.audience, + "headless_callback_url": str(auth_config.headless_callback_url), "success_redirect_url": success_redirect_url, "callback_urls": [str(u) for u in auth_config.callback_urls], } @@ -166,6 +192,7 @@ def _deserialize_auth_config(self, payload: Dict[str, Any]) -> _AuthConfig: token_url=URL(auth_config["token_url"]), client_id=auth_config["client_id"], audience=auth_config["audience"], + headless_callback_url=auth_config["headless_callback_url"], success_redirect_url=success_redirect_url, callback_urls=tuple(URL(u) for u in auth_config.get("callback_urls", [])), ) diff --git a/neuromation/api/login.py b/neuromation/api/login.py index 1ebffb005..19a481c36 100644 --- a/neuromation/api/login.py +++ b/neuromation/api/login.py @@ -112,23 +112,33 @@ async def request(self, code: AuthCode) -> AuthCode: scope="offline_access", audience=self._audience, ) - await self._request(url) + await self._request(url, code) await code.wait() return code @abc.abstractmethod - async def _request(self, url: URL) -> None: + async def _request(self, url: URL, code: AuthCode) -> None: pass class WebBrowserAuthCodeCallbackClient(AuthCodeCallbackClient): - async def _request(self, url: URL) -> None: + async def _request(self, url: URL, code: AuthCode) -> None: loop = asyncio.get_event_loop() await loop.run_in_executor(None, webbrowser.open_new, str(url)) +class HeadlessAuthCodeCallbackClient(AuthCodeCallbackClient): + async def _request(self, url: URL, code: AuthCode) -> None: + print(f"Open {url} in a browser") + print("Put the code displayed in a browser after successful login") + auth_code = input("Code (empty for exit)-> ") + if not auth_code: + sys.exit() + code.set_value(auth_code) + + class DummyAuthCodeCallbackClient(AuthCodeCallbackClient): - async def _request(self, url: URL) -> None: + async def _request(self, url: URL, code: AuthCode) -> None: async with ClientSession() as client: await client.get(url, allow_redirects=True) @@ -341,6 +351,8 @@ class _AuthConfig: client_id: str audience: str + headless_callback_url: URL + callback_urls: Sequence[URL] = ( URL("http://127.0.0.1:54540"), URL("http://127.0.0.1:54541"), @@ -382,16 +394,9 @@ def create( ) -class AuthNegotiator: - def __init__( - self, - config: _AuthConfig, - code_callback_client_factory: Type[ - AuthCodeCallbackClient - ] = WebBrowserAuthCodeCallbackClient, - ) -> None: +class BaseNegotiator: + def __init__(self, config: _AuthConfig) -> None: self._config = config - self._code_callback_client_factory = code_callback_client_factory async def get_code(self) -> AuthCode: code = AuthCode() @@ -422,6 +427,56 @@ async def refresh_token(self, token: Optional[_AuthToken] = None) -> _AuthToken: return token +class AuthNegotiator(BaseNegotiator): + def __init__( + self, + config: _AuthConfig, + code_callback_client_factory: Type[ + AuthCodeCallbackClient + ] = WebBrowserAuthCodeCallbackClient, + ) -> None: + super().__init__(config) + self._code_callback_client_factory = code_callback_client_factory + + async def get_code(self) -> AuthCode: + code = AuthCode() + app = create_auth_code_app(code, redirect_url=self._config.success_redirect_url) + + async with create_app_server( + app, host=self._config.callback_host, ports=self._config.callback_ports + ) as url: + code.callback_url = url + code_callback_client = self._code_callback_client_factory( + url=self._config.auth_url, + client_id=self._config.client_id, + audience=self._config.audience, + ) + return await code_callback_client.request(code) + + +class HeadlessNegotiator(BaseNegotiator): + def __init__( + self, + config: _AuthConfig, + code_callback_client_factory: Type[ + AuthCodeCallbackClient + ] = HeadlessAuthCodeCallbackClient, + ) -> None: + super().__init__(config) + self._code_callback_client_factory = code_callback_client_factory + + async def get_code(self) -> AuthCode: + code = AuthCode() + code.callback_url = self._config.headless_callback_url + + code_callback_client = self._code_callback_client_factory( + url=self._config.auth_url, + client_id=self._config.client_id, + audience=self._config.audience, + ) + return await code_callback_client.request(code) + + @dataclass(frozen=True) class _ServerConfig: auth_config: _AuthConfig @@ -450,6 +505,7 @@ async def get_server_config(url: URL, token: Optional[str] = None) -> _ServerCon if callback_urls is not None else _AuthConfig.callback_urls ) + headless_callback_url = URL(payload["headless_callback_url"]) auth_config = _AuthConfig( auth_url=URL(payload["auth_url"]), token_url=URL(payload["token_url"]), @@ -457,6 +513,7 @@ async def get_server_config(url: URL, token: Optional[str] = None) -> _ServerCon audience=payload["audience"], success_redirect_url=success_redirect_url, callback_urls=callback_urls, + headless_callback_url=headless_callback_url, ) cluster_config = _ClusterConfig( registry_url=URL(payload.get("registry_url", "")), diff --git a/neuromation/cli/config.py b/neuromation/cli/config.py index d89eb8c19..4b2b88d18 100644 --- a/neuromation/cli/config.py +++ b/neuromation/cli/config.py @@ -10,6 +10,7 @@ DEFAULT_API_URL, ConfigError, login as api_login, + login_headless as api_login_headless, login_with_token as api_login_with_token, logout as api_logout, ) @@ -85,6 +86,31 @@ async def login_with_token(root: Root, token: str, url: URL) -> None: click.echo(f"Logged into {url}") +@command() +@click.argument("url", required=False, default=DEFAULT_API_URL, type=URL) +@async_cmd(init_client=False) +async def login_headless(root: Root, url: URL) -> None: + """ + Log into Neuromation Platform from non-GUI server environment. + + URL is a platform entrypoint URL. + + The command works similar to "neuro login" but instead of + opening a browser for performing OAuth registration prints + an URL that should be open on guest host. + + Then user inputs a code displayed in a browser after successful login + back in neuro command to finish the login process. + """ + try: + await api_login_headless(url=url, path=root.config_path, timeout=root.timeout) + except ConfigError: + await api_logout(path=root.config_path) + click.echo("You were successfully logged out.") + await api_login(url=url, path=root.config_path, timeout=root.timeout) + click.echo(f"Logged into {url}") + + @command() @async_cmd() async def logout(root: Root) -> None: @@ -131,6 +157,7 @@ async def docker(root: Root, docker_config: str) -> None: config.add_command(login) config.add_command(login_with_token) +config.add_command(login_headless) config.add_command(show) config.add_command(show_token) From c48a69e63df8752846c13e60aa48dd8ba5d27186 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Thu, 16 May 2019 14:39:58 +0300 Subject: [PATCH 02/28] Fix linter --- neuromation/api/login.py | 26 +++++++++++--------------- tests/api/test_config_factory.py | 3 +++ tests/api/test_login.py | 19 +++++++++++++++++++ tests/api/test_login_utils.py | 8 ++++++++ tests/conftest.py | 1 + 5 files changed, 42 insertions(+), 15 deletions(-) diff --git a/neuromation/api/login.py b/neuromation/api/login.py index 19a481c36..ce0d10d23 100644 --- a/neuromation/api/login.py +++ b/neuromation/api/login.py @@ -4,6 +4,7 @@ import errno import hashlib import secrets +import sys import time import webbrowser from dataclasses import dataclass, field @@ -371,7 +372,11 @@ def callback_ports(self) -> List[int]: def is_initialized(self) -> bool: return bool( - self.auth_url and self.token_url and self.client_id and self.audience + self.auth_url + and self.token_url + and self.client_id + and self.audience + and self.headless_callback_url ) @classmethod @@ -381,6 +386,7 @@ def create( token_url: URL, client_id: str, audience: str, + headless_callback_url: URL, success_redirect_url: Optional[URL] = None, callback_urls: Optional[Sequence[URL]] = None, ) -> "_AuthConfig": @@ -389,29 +395,19 @@ def create( token_url=token_url, client_id=client_id, audience=audience, + headless_callback_url=headless_callback_url, success_redirect_url=success_redirect_url, callback_urls=callback_urls or cls.callback_urls, ) -class BaseNegotiator: +class BaseNegotiator(abc.ABC): def __init__(self, config: _AuthConfig) -> None: self._config = config + @abc.abstractmethod async def get_code(self) -> AuthCode: - code = AuthCode() - app = create_auth_code_app(code, redirect_url=self._config.success_redirect_url) - - async with create_app_server( - app, host=self._config.callback_host, ports=self._config.callback_ports - ) as url: - code.callback_url = url - code_callback_client = self._code_callback_client_factory( - url=self._config.auth_url, - client_id=self._config.client_id, - audience=self._config.audience, - ) - return await code_callback_client.request(code) + pass async def refresh_token(self, token: Optional[_AuthToken] = None) -> _AuthToken: async with AuthTokenClient( diff --git a/tests/api/test_config_factory.py b/tests/api/test_config_factory.py index 1bdcd8c14..ec630dccd 100644 --- a/tests/api/test_config_factory.py +++ b/tests/api/test_config_factory.py @@ -107,6 +107,7 @@ def test_check_initialized(self) -> None: token_url=URL("http://token"), client_id="client-id", audience="everyone", + headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), ) assert auth_config_good.is_initialized() @@ -135,6 +136,7 @@ def test_check_initialized_bad_auth_config(self) -> None: token_url=URL("http://token"), client_id="client-id", audience="everyone", + headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), ) assert not auth_config_bad.is_initialized() @@ -164,6 +166,7 @@ def test_check_initialized_bad_cluster_config(self) -> None: token_url=URL("http://token"), client_id="client-id", audience="everyone", + headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), ) assert auth_config_bad.is_initialized() diff --git a/tests/api/test_login.py b/tests/api/test_login.py index ec1914d77..18f80724c 100644 --- a/tests/api/test_login.py +++ b/tests/api/test_login.py @@ -315,6 +315,7 @@ async def auth_config( token_url=auth_server / "oauth/token", client_id=auth_client_id, audience="https://platform.dev.neuromation.io", + headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), callback_urls=[URL(f"http://127.0.0.1:{port}")], ) @@ -386,6 +387,7 @@ def test_is_initialized__no_auth_url(self) -> None: token_url=URL("url"), client_id="client_id", audience="audience", + headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), callback_urls=(URL("url1"), URL("url2")), success_redirect_url=URL("url"), ) @@ -397,6 +399,7 @@ def test_is_initialized__no_token_url(self) -> None: token_url=URL(), client_id="client_id", audience="audience", + headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), callback_urls=(URL("url1"), URL("url2")), success_redirect_url=URL("url"), ) @@ -408,6 +411,7 @@ def test_is_initialized__no_client_id(self) -> None: token_url=URL("url"), client_id="", audience="audience", + headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), callback_urls=(URL("url1"), URL("url2")), success_redirect_url=URL("url"), ) @@ -419,6 +423,7 @@ def test_is_initialized__no_audience(self) -> None: token_url=URL("url"), client_id="client_id", audience="", + headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), callback_urls=(URL("url1"), URL("url2")), success_redirect_url=URL("url"), ) @@ -430,6 +435,7 @@ def test_is_initialized__no_callback_urls(self) -> None: token_url=URL("url"), client_id="client_id", audience="audience", + headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), callback_urls=[], success_redirect_url=URL("url"), ) @@ -441,11 +447,24 @@ def test_is_initialized__no_success_redirect_url(self) -> None: token_url=URL("url"), client_id="client_id", audience="audience", + headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), callback_urls=(URL("url1"), URL("url2")), success_redirect_url=None, ) assert auth_config.is_initialized() is True + def test_is_initialized__no_headless_callback_url(self) -> None: + auth_config = _AuthConfig( + auth_url=URL("url"), + token_url=URL("url"), + client_id="client_id", + audience="audience", + headless_callback_url=URL(), + callback_urls=(URL("url1"), URL("url2")), + success_redirect_url=None, + ) + assert auth_config.is_initialized() is False + class TestClusterConfig: def test_is_initialized(self) -> None: diff --git a/tests/api/test_login_utils.py b/tests/api/test_login_utils.py index b3871b8bc..e236bce56 100644 --- a/tests/api/test_login_utils.py +++ b/tests/api/test_login_utils.py @@ -17,6 +17,7 @@ async def test_get_server_config(aiohttp_server: _TestServerFactory) -> None: token_url = "https://dev-neuromation.auth0.com/oauth/token" client_id = "this_is_client_id" audience = "https://platform.dev.neuromation.io" + headless_callback_url = "https://https://dev.neu.ro/oauth/show-code" callback_urls = [ "http://127.0.0.1:54540", "http://127.0.0.1:54541", @@ -30,6 +31,7 @@ async def test_get_server_config(aiohttp_server: _TestServerFactory) -> None: "audience": audience, "callback_urls": callback_urls, "success_redirect_url": success_redirect_url, + "headless_callback_url": headless_callback_url, } async def handler(request: web.Request) -> web.Response: @@ -47,6 +49,7 @@ async def handler(request: web.Request) -> web.Response: token_url=URL(token_url), client_id=client_id, audience=audience, + headless_callback_url=URL(headless_callback_url), callback_urls=tuple(URL(u) for u in callback_urls), success_redirect_url=URL(success_redirect_url), ), @@ -63,6 +66,7 @@ async def test_get_server_config_no_callback_urls( token_url = "https://dev-neuromation.auth0.com/oauth/token" client_id = "this_is_client_id" audience = "https://platform.dev.neuromation.io" + headless_callback_url = "https://https://dev.neu.ro/oauth/show-code" success_redirect_url = "https://platform.neuromation.io" JSON = { "auth_url": auth_url, @@ -87,6 +91,7 @@ async def handler(request: web.Request) -> web.Response: token_url=URL(token_url), client_id=client_id, audience=audience, + headless_callback_url=URL(headless_callback_url), success_redirect_url=URL(success_redirect_url), ), cluster_config=_ClusterConfig( @@ -104,6 +109,7 @@ async def test_get_server_config_with_token(aiohttp_server: _TestServerFactory) token_url = "https://dev-neuromation.auth0.com/oauth/token" client_id = "this_is_client_id" audience = "https://platform.dev.neuromation.io" + headless_callback_url = "https://https://dev.neu.ro/oauth/show-code" success_redirect_url = "https://platform.neuromation.io" JSON = { "registry_url": registry_url, @@ -114,6 +120,7 @@ async def test_get_server_config_with_token(aiohttp_server: _TestServerFactory) "token_url": token_url, "client_id": client_id, "audience": audience, + "headless_callback_url": headless_callback_url, "success_redirect_url": success_redirect_url, } @@ -132,6 +139,7 @@ async def handler(request: web.Request) -> web.Response: token_url=URL(token_url), client_id=client_id, audience=audience, + headless_callback_url=URL(headless_callback_url), success_redirect_url=URL(success_redirect_url), ), cluster_config=_ClusterConfig.create( diff --git a/tests/conftest.py b/tests/conftest.py index 4abd12494..e07ffc96f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -21,6 +21,7 @@ def auth_config() -> _AuthConfig: token_url=URL("https://dev-neuromation.auth0.com/oauth/token"), client_id="CLIENT-ID", audience="https://platform.dev.neuromation.io", + headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), success_redirect_url=URL("https://neu.ro/#running-your-first-job"), callback_urls=[ URL("http://127.0.0.1:54540"), From bd513bcbe63f04fa4791b08add8bc54179a65424 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Thu, 16 May 2019 14:48:31 +0300 Subject: [PATCH 03/28] Fix tests --- tests/api/test_config_factory.py | 1 + tests/api/test_login_utils.py | 1 + 2 files changed, 2 insertions(+) diff --git a/tests/api/test_config_factory.py b/tests/api/test_config_factory.py index ec630dccd..65084f5ad 100644 --- a/tests/api/test_config_factory.py +++ b/tests/api/test_config_factory.py @@ -55,6 +55,7 @@ async def _config_handler(request: web.Request) -> web.Response: "token_url": "https://test-neuromation.auth0.com/oauth/token", "client_id": "banana", "audience": "https://test.dev.neuromation.io", + "headless_callback_url": "https://https://dev.neu.ro/oauth/show-code", "callback_urls": [ "http://127.0.0.2:54540", "http://127.0.0.2:54541", diff --git a/tests/api/test_login_utils.py b/tests/api/test_login_utils.py index e236bce56..e3e25dcf6 100644 --- a/tests/api/test_login_utils.py +++ b/tests/api/test_login_utils.py @@ -73,6 +73,7 @@ async def test_get_server_config_no_callback_urls( "token_url": token_url, "client_id": client_id, "audience": audience, + "headless_callback_url": headless_callback_url, "success_redirect_url": success_redirect_url, } From 6603724aed4e5a49f5f33d062cd90f30b3c07bd8 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Fri, 17 May 2019 12:17:33 +0300 Subject: [PATCH 04/28] Inline method --- neuromation/api/config_factory.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/neuromation/api/config_factory.py b/neuromation/api/config_factory.py index 0abca3eaa..dbb4b3b03 100644 --- a/neuromation/api/config_factory.py +++ b/neuromation/api/config_factory.py @@ -39,7 +39,8 @@ def __init__(self, path: Optional[Path] = None) -> None: async def get(self, *, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT) -> Client: config = self._read() - new_token = await self._refresh_auth_token(config) + auth_negotiator = AuthNegotiator(config=config.auth_config) + new_token = await auth_negotiator.refresh_token(config.auth_token) if new_token != config.auth_token: new_config = replace(config, auth_token=new_token) self._save(new_config) @@ -210,10 +211,6 @@ def _deserialize_auth_token(self, payload: Dict[str, Any]) -> _AuthToken: refresh_token=auth_payload["refresh_token"], ) - async def _refresh_auth_token(self, config: _Config) -> _AuthToken: - auth_negotiator = AuthNegotiator(config=config.auth_config) - return await auth_negotiator.refresh_token(config.auth_token) - def _save(self, config: _Config) -> None: payload: Dict[str, Any] = dict() payload["url"] = str(config.url) From c7d72ceaa717493c952230efc388919882ef799f Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Fri, 17 May 2019 12:26:27 +0300 Subject: [PATCH 05/28] Make token refreshing explicit --- neuromation/api/config_factory.py | 4 ++-- neuromation/api/login.py | 9 +++++++++ tests/api/test_config_factory.py | 9 +++++++-- 3 files changed, 18 insertions(+), 4 deletions(-) diff --git a/neuromation/api/config_factory.py b/neuromation/api/config_factory.py index dbb4b3b03..8a42266fb 100644 --- a/neuromation/api/config_factory.py +++ b/neuromation/api/config_factory.py @@ -18,6 +18,7 @@ _AuthToken, _ClusterConfig, get_server_config, + refresh_token, ) @@ -39,8 +40,7 @@ def __init__(self, path: Optional[Path] = None) -> None: async def get(self, *, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT) -> Client: config = self._read() - auth_negotiator = AuthNegotiator(config=config.auth_config) - new_token = await auth_negotiator.refresh_token(config.auth_token) + new_token = await refresh_token(config.auth_config, config.auth_token) if new_token != config.auth_token: new_config = replace(config, auth_token=new_token) self._save(new_config) diff --git a/neuromation/api/login.py b/neuromation/api/login.py index ce0d10d23..72b4cc71f 100644 --- a/neuromation/api/login.py +++ b/neuromation/api/login.py @@ -401,6 +401,15 @@ def create( ) +async def refresh_token(config: _AuthConfig, token: _AuthToken) -> _AuthToken: + async with AuthTokenClient( + url=config.token_url, client_id=config.client_id + ) as token_client: + if token.is_expired: + return await token_client.refresh(token) + return token + + class BaseNegotiator(abc.ABC): def __init__(self, config: _AuthConfig) -> None: self._config = config diff --git a/tests/api/test_config_factory.py b/tests/api/test_config_factory.py index 65084f5ad..d531004eb 100644 --- a/tests/api/test_config_factory.py +++ b/tests/api/test_config_factory.py @@ -8,6 +8,7 @@ from aiohttp import web from yarl import URL +import neuromation.api.config_factory from neuromation.api import ConfigError, Factory from neuromation.api.config import _AuthConfig, _AuthToken, _Config, _PyPIVersion from neuromation.api.jobs import Jobs @@ -232,10 +233,14 @@ async def test_token_autorefreshing( ) -> None: new_token = str(uuid()) + "changed" * 10 # token must has other size - async def _refresh_token_mock(self: Any, token: str) -> _AuthToken: + async def _refresh_token_mock( + configf: _AuthConfig, token: _AuthToken + ) -> _AuthToken: return _AuthToken.create_non_expiring(new_token) - monkeypatch.setattr(AuthNegotiator, "refresh_token", _refresh_token_mock) + monkeypatch.setattr( + neuromation.api.config_factory, "refresh_token", _refresh_token_mock + ) file_stat_before = config_file.stat() client = await Factory().get() await client.close() From 722916771d79569d5a1728697dfa4fbd72a3afcc Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Fri, 17 May 2019 15:17:01 +0300 Subject: [PATCH 06/28] Fix e2e by using real login call instead of constant config file --- tests/e2e/conftest.py | 6 +----- tests/e2e/utils.py | 25 ------------------------- 2 files changed, 1 insertion(+), 30 deletions(-) diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 42f4f2916..08c3de76a 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -44,7 +44,6 @@ FILE_SIZE_B, JOB_TINY_CONTAINER_PARAMS, NGINX_IMAGE_NAME, - RC_TEXT, JobWaitStateStopReached, ) @@ -512,10 +511,7 @@ def nmrc_path(tmp_path: Path) -> Optional[Path]: e2e_test_token = os.environ.get("CLIENT_TEST_E2E_USER_NAME") if e2e_test_token: nmrc_path = tmp_path / "conftest.nmrc" - - rc_text = RC_TEXT.format(token=e2e_test_token) - nmrc_path.write_text(rc_text) - nmrc_path.chmod(0o600) + run(login_with_token(e2e_test_token, "https://dev.neu.ro/api/v1")) return nmrc_path else: return None diff --git a/tests/e2e/utils.py b/tests/e2e/utils.py index 2979e6c21..5b620b136 100644 --- a/tests/e2e/utils.py +++ b/tests/e2e/utils.py @@ -2,31 +2,6 @@ FILE_SIZE_MB = 16 FILE_SIZE_B = FILE_SIZE_MB * 1024 * 1024 GENERATION_TIMEOUT_SEC = 120 -RC_TEXT = """ -auth_config: - audience: https://platform.dev.neuromation.io - auth_url: https://dev-neuromation.auth0.com/authorize - callback_urls: - - http://127.0.0.1:54540 - - http://127.0.0.1:54541 - - http://127.0.0.1:54542 - client_id: CLIENT-ID - success_redirect_url: https://neu.ro/#running-your-first-job - token_url: https://dev-neuromation.auth0.com/oauth/token -auth_token: - expiration_time: 1713014496 - refresh_token: refresh-token - token: {token} -pypi: - check_timestamp: 0 - pypi_version: 0.0.0 -cluster_config: - monitoring_url: https://dev.neu.ro/api/v1/jobs - registry_url: https://registry-dev.neu.ro - storage_url: https://dev.neu.ro/api/v1/storage - users_url: https://dev.neu.ro/api/v1/users -url: https://dev.neu.ro/api/v1 -""" UBUNTU_IMAGE_NAME = "ubuntu:latest" NGINX_IMAGE_NAME = "nginx:latest" From 7a9ddf76534b115dddd45c9029adae4871033070 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Fri, 17 May 2019 15:23:42 +0300 Subject: [PATCH 07/28] Fix import --- tests/e2e/conftest.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 08c3de76a..94803918c 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -36,6 +36,7 @@ JobStatus, ResourceNotFound, get as api_get, + login_with_token, ) from neuromation.cli import main from neuromation.cli.const import EX_IOERR, EX_OK, EX_OSFILE From 7802bbc2cfc91f82f3f315daab0788bffe008b5a Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Fri, 17 May 2019 15:24:55 +0300 Subject: [PATCH 08/28] Fix test --- tests/e2e/conftest.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 94803918c..8502a625d 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -512,7 +512,11 @@ def nmrc_path(tmp_path: Path) -> Optional[Path]: e2e_test_token = os.environ.get("CLIENT_TEST_E2E_USER_NAME") if e2e_test_token: nmrc_path = tmp_path / "conftest.nmrc" - run(login_with_token(e2e_test_token, "https://dev.neu.ro/api/v1")) + run( + login_with_token( + e2e_test_token, url=URL("https://dev.neu.ro/api/v1"), path=nmrc_path + ) + ) return nmrc_path else: return None From c612c8662dcd6a0fbb30bc6fdbbfe5c55b2ef5e9 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Fri, 17 May 2019 15:31:32 +0300 Subject: [PATCH 09/28] Use session scope for e2e login --- tests/e2e/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 8502a625d..56d899795 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -507,7 +507,7 @@ def get_last_output(self) -> SysCap: return self._last_output -@pytest.fixture() +@pytest.fixture(scope="session") def nmrc_path(tmp_path: Path) -> Optional[Path]: e2e_test_token = os.environ.get("CLIENT_TEST_E2E_USER_NAME") if e2e_test_token: From 58e205b96037f01f012f3d2afcad05c183492cd1 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Fri, 17 May 2019 15:33:24 +0300 Subject: [PATCH 10/28] Fix test --- tests/e2e/conftest.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 56d899795..8085559d3 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -508,9 +508,10 @@ def get_last_output(self) -> SysCap: @pytest.fixture(scope="session") -def nmrc_path(tmp_path: Path) -> Optional[Path]: +def nmrc_path(tmp_path_factory: Any) -> Optional[Path]: e2e_test_token = os.environ.get("CLIENT_TEST_E2E_USER_NAME") if e2e_test_token: + tmp_path = tmp_path_factory.mktemp("config") nmrc_path = tmp_path / "conftest.nmrc" run( login_with_token( From f1bafa707b6ccc52890180ee56cecd7c2c434e4c Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Fri, 17 May 2019 16:43:43 +0300 Subject: [PATCH 11/28] Pass timeout everywhere (almost) --- neuromation/api/config_factory.py | 6 +++--- neuromation/api/login.py | 25 +++++++++++++++++-------- tests/api/test_config_factory.py | 3 ++- tests/api/test_login.py | 25 ++++++++++++++++++------- 4 files changed, 40 insertions(+), 19 deletions(-) diff --git a/neuromation/api/config_factory.py b/neuromation/api/config_factory.py index 8a42266fb..058c3d931 100644 --- a/neuromation/api/config_factory.py +++ b/neuromation/api/config_factory.py @@ -40,7 +40,7 @@ def __init__(self, path: Optional[Path] = None) -> None: async def get(self, *, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT) -> Client: config = self._read() - new_token = await refresh_token(config.auth_config, config.auth_token) + new_token = await refresh_token(config.auth_config, config.auth_token, timeout) if new_token != config.auth_token: new_config = replace(config, auth_token=new_token) self._save(new_config) @@ -56,7 +56,7 @@ async def login( if self._path.exists(): raise ConfigError(f"Config file {self._path} already exists. Please logout") config_unauthorized = await get_server_config(url) - negotiator = AuthNegotiator(config_unauthorized.auth_config) + negotiator = AuthNegotiator(config_unauthorized.auth_config, timeout) auth_token = await negotiator.refresh_token() config_authorized = await get_server_config(url, token=auth_token.token) @@ -80,7 +80,7 @@ async def login_headless( if self._path.exists(): raise ConfigError(f"Config file {self._path} already exists. Please logout") config_unauthorized = await get_server_config(url) - negotiator = HeadlessNegotiator(config_unauthorized.auth_config) + negotiator = HeadlessNegotiator(config_unauthorized.auth_config, timeout) auth_token = await negotiator.refresh_token() config_authorized = await get_server_config(url, token=auth_token.token) diff --git a/neuromation/api/login.py b/neuromation/api/login.py index 72b4cc71f..9137724e5 100644 --- a/neuromation/api/login.py +++ b/neuromation/api/login.py @@ -268,11 +268,13 @@ def create_non_expiring(cls, token: str) -> "_AuthToken": class AuthTokenClient: - def __init__(self, url: URL, client_id: str) -> None: + def __init__( + self, url: URL, client_id: str, timeout: aiohttp.ClientTimeout + ) -> None: self._url = url self._client_id = client_id - self._client = ClientSession() + self._client = ClientSession(timeout=timeout) async def close(self) -> None: await self._client.close() @@ -401,9 +403,11 @@ def create( ) -async def refresh_token(config: _AuthConfig, token: _AuthToken) -> _AuthToken: +async def refresh_token( + config: _AuthConfig, token: _AuthToken, timeout: aiohttp.ClientTimeout +) -> _AuthToken: async with AuthTokenClient( - url=config.token_url, client_id=config.client_id + url=config.token_url, client_id=config.client_id, timeout=timeout ) as token_client: if token.is_expired: return await token_client.refresh(token) @@ -411,8 +415,9 @@ async def refresh_token(config: _AuthConfig, token: _AuthToken) -> _AuthToken: class BaseNegotiator(abc.ABC): - def __init__(self, config: _AuthConfig) -> None: + def __init__(self, config: _AuthConfig, timeout: aiohttp.ClientTimeout) -> None: self._config = config + self._timeout = timeout @abc.abstractmethod async def get_code(self) -> AuthCode: @@ -420,7 +425,9 @@ async def get_code(self) -> AuthCode: async def refresh_token(self, token: Optional[_AuthToken] = None) -> _AuthToken: async with AuthTokenClient( - url=self._config.token_url, client_id=self._config.client_id + url=self._config.token_url, + client_id=self._config.client_id, + timeout=self._timeout, ) as token_client: if not token: code = await self.get_code() @@ -436,11 +443,12 @@ class AuthNegotiator(BaseNegotiator): def __init__( self, config: _AuthConfig, + timeout: aiohttp.ClientTimeout, code_callback_client_factory: Type[ AuthCodeCallbackClient ] = WebBrowserAuthCodeCallbackClient, ) -> None: - super().__init__(config) + super().__init__(config, timeout) self._code_callback_client_factory = code_callback_client_factory async def get_code(self) -> AuthCode: @@ -463,11 +471,12 @@ class HeadlessNegotiator(BaseNegotiator): def __init__( self, config: _AuthConfig, + timeout: aiohttp.ClientTimeout, code_callback_client_factory: Type[ AuthCodeCallbackClient ] = HeadlessAuthCodeCallbackClient, ) -> None: - super().__init__(config) + super().__init__(config, timeout) self._code_callback_client_factory = code_callback_client_factory async def get_code(self) -> AuthCode: diff --git a/tests/api/test_config_factory.py b/tests/api/test_config_factory.py index d531004eb..fd2bc67c7 100644 --- a/tests/api/test_config_factory.py +++ b/tests/api/test_config_factory.py @@ -3,6 +3,7 @@ from typing import Any, List, Optional, Set from uuid import uuid4 as uuid +import aiohttp import pytest import yaml from aiohttp import web @@ -234,7 +235,7 @@ async def test_token_autorefreshing( new_token = str(uuid()) + "changed" * 10 # token must has other size async def _refresh_token_mock( - configf: _AuthConfig, token: _AuthToken + configf: _AuthConfig, token: _AuthToken, timeout: aiohttp.ClientTimeout ) -> _AuthToken: return _AuthToken.create_non_expiring(new_token) diff --git a/tests/api/test_login.py b/tests/api/test_login.py index 18f80724c..db9b5de43 100644 --- a/tests/api/test_login.py +++ b/tests/api/test_login.py @@ -18,6 +18,7 @@ ) from yarl import URL +from neuromation.api.core import DEFAULT_TIMEOUT from neuromation.api.login import ( AuthCode, AuthException, @@ -327,7 +328,7 @@ async def test_request(self, auth_client_id: str, auth_config: _AuthConfig) -> N code.callback_url = auth_config.callback_urls[0] async with AuthTokenClient( - auth_config.token_url, client_id=auth_client_id + auth_config.token_url, client_id=auth_client_id, timeout=DEFAULT_TIMEOUT ) as client: token = await client.request(code) assert token.token == "test_access_token" @@ -342,7 +343,7 @@ async def test_refresh(self, auth_client_id: str, auth_config: _AuthConfig) -> N ) async with AuthTokenClient( - auth_config.token_url, client_id=auth_client_id + auth_config.token_url, client_id=auth_client_id, timeout=DEFAULT_TIMEOUT ) as client: new_token = await client.refresh(token) assert new_token.token == "test_access_token_refreshed" @@ -367,7 +368,9 @@ async def handle_token(request: Request) -> Response: server = await aiohttp_server(app) url = server.make_url("/oauth/token") - async with AuthTokenClient(url, client_id=client_id) as client: + async with AuthTokenClient( + url, client_id=client_id, timeout=DEFAULT_TIMEOUT + ) as client: with pytest.raises(AuthException, match="failed to get an access token."): await client.request(code) @@ -516,7 +519,9 @@ def test_is_initialized__no_monitoring_url(self) -> None: class TestAuthNegotiator: async def test_get_code(self, auth_config: _AuthConfig) -> None: negotiator = AuthNegotiator( - config=auth_config, code_callback_client_factory=DummyAuthCodeCallbackClient + config=auth_config, + timeout=DEFAULT_TIMEOUT, + code_callback_client_factory=DummyAuthCodeCallbackClient, ) code = await negotiator.get_code() assert await code.wait() == "test_code" @@ -524,7 +529,9 @@ async def test_get_code(self, auth_config: _AuthConfig) -> None: async def test_get_token(self, auth_config: _AuthConfig) -> None: negotiator = AuthNegotiator( - config=auth_config, code_callback_client_factory=DummyAuthCodeCallbackClient + config=auth_config, + timeout=DEFAULT_TIMEOUT, + code_callback_client_factory=DummyAuthCodeCallbackClient, ) token = await negotiator.refresh_token(token=None) assert token.token == "test_access_token" @@ -532,7 +539,9 @@ async def test_get_token(self, auth_config: _AuthConfig) -> None: async def test_refresh_token_noop(self, auth_config: _AuthConfig) -> None: negotiator = AuthNegotiator( - config=auth_config, code_callback_client_factory=DummyAuthCodeCallbackClient + config=auth_config, + timeout=DEFAULT_TIMEOUT, + code_callback_client_factory=DummyAuthCodeCallbackClient, ) token = await negotiator.refresh_token(token=None) assert token.token == "test_access_token" @@ -545,7 +554,9 @@ async def test_refresh_token_noop(self, auth_config: _AuthConfig) -> None: async def test_refresh_token(self, auth_config: _AuthConfig) -> None: negotiator = AuthNegotiator( - config=auth_config, code_callback_client_factory=DummyAuthCodeCallbackClient + config=auth_config, + timeout=DEFAULT_TIMEOUT, + code_callback_client_factory=DummyAuthCodeCallbackClient, ) token = await negotiator.refresh_token(token=None) assert token.token == "test_access_token" From 00f0db35a150227e7d4c86b32d6fa9e698edccd6 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Fri, 17 May 2019 16:48:50 +0300 Subject: [PATCH 12/28] Increase login timeout in tests --- tests/e2e/conftest.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 8085559d3..de0b4b99f 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -515,7 +515,10 @@ def nmrc_path(tmp_path_factory: Any) -> Optional[Path]: nmrc_path = tmp_path / "conftest.nmrc" run( login_with_token( - e2e_test_token, url=URL("https://dev.neu.ro/api/v1"), path=nmrc_path + e2e_test_token, + url=URL("https://dev.neu.ro/api/v1"), + path=nmrc_path, + timeout=CLIENT_TIMEOUT, ) ) return nmrc_path From ba23a35246bf59ee3ba1c8315357fed37299be24 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 20 May 2019 12:59:01 +0300 Subject: [PATCH 13/28] Refactor headless login api --- neuromation/api/__init__.py | 3 ++- neuromation/api/config_factory.py | 7 +++++-- neuromation/api/login.py | 27 +++++++++++++++++---------- neuromation/cli/config.py | 15 ++++++++++++++- 4 files changed, 38 insertions(+), 14 deletions(-) diff --git a/neuromation/api/__init__.py b/neuromation/api/__init__.py index 422ec04f4..979bccffb 100644 --- a/neuromation/api/__init__.py +++ b/neuromation/api/__init__.py @@ -1,7 +1,7 @@ import sys from pathlib import Path from types import TracebackType -from typing import Any, Awaitable, Coroutine, Generator, Optional, Type +from typing import Any, Awaitable, Callable, Coroutine, Generator, Optional, Type import aiohttp from yarl import URL @@ -160,6 +160,7 @@ async def login_with_token( async def login_headless( + callback: Callable[[URL], Awaitable[str]], *, url: URL = DEFAULT_API_URL, path: Optional[Path] = None, diff --git a/neuromation/api/config_factory.py b/neuromation/api/config_factory.py index 058c3d931..581e25ad8 100644 --- a/neuromation/api/config_factory.py +++ b/neuromation/api/config_factory.py @@ -2,7 +2,7 @@ import sys from dataclasses import replace from pathlib import Path -from typing import Any, Dict, Optional +from typing import Any, Awaitable, Callable, Dict, Optional import aiohttp import yaml @@ -73,6 +73,7 @@ async def login( async def login_headless( self, + callback: Callable[[URL], Awaitable[str]], *, url: URL = DEFAULT_API_URL, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT, @@ -80,7 +81,9 @@ async def login_headless( if self._path.exists(): raise ConfigError(f"Config file {self._path} already exists. Please logout") config_unauthorized = await get_server_config(url) - negotiator = HeadlessNegotiator(config_unauthorized.auth_config, timeout) + negotiator = HeadlessNegotiator( + config_unauthorized.auth_config, callback, timeout + ) auth_token = await negotiator.refresh_token() config_authorized = await get_server_config(url, token=auth_token.token) diff --git a/neuromation/api/login.py b/neuromation/api/login.py index 9137724e5..e56dc9657 100644 --- a/neuromation/api/login.py +++ b/neuromation/api/login.py @@ -11,6 +11,7 @@ from typing import ( Any, AsyncIterator, + Awaitable, Callable, Dict, List, @@ -129,12 +130,19 @@ async def _request(self, url: URL, code: AuthCode) -> None: class HeadlessAuthCodeCallbackClient(AuthCodeCallbackClient): + def __init__( + self, + url: URL, + client_id: str, + audience: str, + callback: Callable[[URL], Awaitable[str]], + ) -> None: + super().__init__(url=url, client_id=client_id, audience=audience) + self._callback = callback + async def _request(self, url: URL, code: AuthCode) -> None: - print(f"Open {url} in a browser") - print("Put the code displayed in a browser after successful login") - auth_code = input("Code (empty for exit)-> ") - if not auth_code: - sys.exit() + auth_code = await self._callback(url) + assert auth_code code.set_value(auth_code) @@ -471,22 +479,21 @@ class HeadlessNegotiator(BaseNegotiator): def __init__( self, config: _AuthConfig, + callback: Callable[[URL], Awaitable[str]], timeout: aiohttp.ClientTimeout, - code_callback_client_factory: Type[ - AuthCodeCallbackClient - ] = HeadlessAuthCodeCallbackClient, ) -> None: super().__init__(config, timeout) - self._code_callback_client_factory = code_callback_client_factory + self._callback = callback async def get_code(self) -> AuthCode: code = AuthCode() code.callback_url = self._config.headless_callback_url - code_callback_client = self._code_callback_client_factory( + code_callback_client = AuthCodeCallbackClient( url=self._config.auth_url, client_id=self._config.client_id, audience=self._config.audience, + callback=self._callback, ) return await code_callback_client.request(code) diff --git a/neuromation/cli/config.py b/neuromation/cli/config.py index a9f795002..47bd440b9 100644 --- a/neuromation/cli/config.py +++ b/neuromation/cli/config.py @@ -102,8 +102,21 @@ async def login_headless(root: Root, url: URL) -> None: Then user inputs a code displayed in a browser after successful login back in neuro command to finish the login process. """ + + async def login_callback(url: URL) -> str: + click.echo(f"Open {url} in a browser") + click.echo("Put the code displayed in a browser after successful login") + auth_code = input("Code (empty for exit)-> ") + if not auth_code: + sys.exit() + try: - await api_login_headless(url=url, path=root.config_path, timeout=root.timeout) + await api_login_headless( + url=url, + callback=login_callback, + path=root.config_path, + timeout=root.timeout, + ) except ConfigError: await api_logout(path=root.config_path) click.echo("You were successfully logged out.") From 7edc836ae4bfb3f2a0eb4d5fc776efcab04383f5 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 20 May 2019 13:05:01 +0300 Subject: [PATCH 14/28] Refactor headless login api --- neuromation/api/login.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/neuromation/api/login.py b/neuromation/api/login.py index e56dc9657..6f2ef255f 100644 --- a/neuromation/api/login.py +++ b/neuromation/api/login.py @@ -489,7 +489,7 @@ async def get_code(self) -> AuthCode: code = AuthCode() code.callback_url = self._config.headless_callback_url - code_callback_client = AuthCodeCallbackClient( + code_callback_client = HeadlessAuthCodeCallbackClient( url=self._config.auth_url, client_id=self._config.client_id, audience=self._config.audience, From 5ec45fdab2b80a0affe2022aad96b8cd2820fde7 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 20 May 2019 13:07:13 +0300 Subject: [PATCH 15/28] Fix typo --- neuromation/api/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/neuromation/api/__init__.py b/neuromation/api/__init__.py index 979bccffb..898fd6fde 100644 --- a/neuromation/api/__init__.py +++ b/neuromation/api/__init__.py @@ -166,7 +166,7 @@ async def login_headless( path: Optional[Path] = None, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT ) -> None: - await Factory(path).login_headless(url=url, timeout=timeout) + await Factory(path).login_headless(callback, url=url, timeout=timeout) async def logout(*, path: Optional[Path] = None) -> None: From b634786c196e50e550144923de4b916ea7708192 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 20 May 2019 13:08:52 +0300 Subject: [PATCH 16/28] Fix imports --- neuromation/api/login.py | 1 - neuromation/cli/config.py | 2 ++ 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/neuromation/api/login.py b/neuromation/api/login.py index 6f2ef255f..f5498f4fc 100644 --- a/neuromation/api/login.py +++ b/neuromation/api/login.py @@ -4,7 +4,6 @@ import errno import hashlib import secrets -import sys import time import webbrowser from dataclasses import dataclass, field diff --git a/neuromation/cli/config.py b/neuromation/cli/config.py index 47bd440b9..dec804f2f 100644 --- a/neuromation/cli/config.py +++ b/neuromation/cli/config.py @@ -1,5 +1,6 @@ import json import os +import sys from pathlib import Path from typing import Any, Dict @@ -109,6 +110,7 @@ async def login_callback(url: URL) -> str: auth_code = input("Code (empty for exit)-> ") if not auth_code: sys.exit() + return auth_code try: await api_login_headless( From 9348872d2e3e97c7f70188842a1ac5724865071b Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 20 May 2019 13:37:55 +0300 Subject: [PATCH 17/28] Add test --- tests/api/test_login.py | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/tests/api/test_login.py b/tests/api/test_login.py index db9b5de43..bafd8f041 100644 --- a/tests/api/test_login.py +++ b/tests/api/test_login.py @@ -25,6 +25,7 @@ AuthNegotiator, AuthTokenClient, DummyAuthCodeCallbackClient, + HeadlessNegotiator, _AuthConfig, _AuthToken, _ClusterConfig, @@ -316,7 +317,7 @@ async def auth_config( token_url=auth_server / "oauth/token", client_id=auth_client_id, audience="https://platform.dev.neuromation.io", - headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), + headless_callback_url=URL("https://dev.neu.ro/oauth/show-code"), callback_urls=[URL(f"http://127.0.0.1:{port}")], ) @@ -569,3 +570,25 @@ async def test_refresh_token(self, auth_config: _AuthConfig) -> None: token = await negotiator.refresh_token(token=token) assert token.token == "test_access_token_refreshed" assert token.refresh_token == "test_refresh_token" + + +class TestHeadlessNegotiator: + async def test_get_code(self, auth_config: _AuthConfig) -> None: + async def callback(url: URL) -> str: + assert url.with_query(None) == auth_config.auth_url + + assert dict(url.query) == dict( + response_type='code', + code_challenge=mock.ANY, + code_challenge_method='S256', + client_id='test_client_id', + redirect_uri='https://dev.neu.ro/oauth/show-code', + scope='offline_access', + audience="https://platform.dev.neuromation.io") + return "test_code" + + negotiator = HeadlessNegotiator( + config=auth_config, callback=callback, timeout=DEFAULT_TIMEOUT + ) + code = await negotiator.get_code() + assert await code.wait() == "test_code" From 5cca24bf38739f592677872cf85a25a50d1ae927 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 20 May 2019 13:38:11 +0300 Subject: [PATCH 18/28] Fix linter --- tests/api/test_login.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/api/test_login.py b/tests/api/test_login.py index bafd8f041..8df499782 100644 --- a/tests/api/test_login.py +++ b/tests/api/test_login.py @@ -578,13 +578,14 @@ async def callback(url: URL) -> str: assert url.with_query(None) == auth_config.auth_url assert dict(url.query) == dict( - response_type='code', + response_type="code", code_challenge=mock.ANY, - code_challenge_method='S256', - client_id='test_client_id', - redirect_uri='https://dev.neu.ro/oauth/show-code', - scope='offline_access', - audience="https://platform.dev.neuromation.io") + code_challenge_method="S256", + client_id="test_client_id", + redirect_uri="https://dev.neu.ro/oauth/show-code", + scope="offline_access", + audience="https://platform.dev.neuromation.io", + ) return "test_code" negotiator = HeadlessNegotiator( From f8c3f8bc2fd72f1255768bfd8f64d32526dbf45f Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 20 May 2019 13:41:07 +0300 Subject: [PATCH 19/28] Another test --- tests/api/test_login.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/api/test_login.py b/tests/api/test_login.py index 8df499782..ef1292584 100644 --- a/tests/api/test_login.py +++ b/tests/api/test_login.py @@ -593,3 +593,13 @@ async def callback(url: URL) -> str: ) code = await negotiator.get_code() assert await code.wait() == "test_code" + + async def test_get_code_raises(self, auth_config: _AuthConfig) -> None: + async def callback(url: URL) -> str: + raise RuntimeError("callback error") + + negotiator = HeadlessNegotiator( + config=auth_config, callback=callback, timeout=DEFAULT_TIMEOUT + ) + with pytest.raises(RuntimeError, match="callback error"): + await negotiator.get_code() From 83edc47a997419a29eb3e390a9d63418ec669575 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 20 May 2019 13:44:52 +0300 Subject: [PATCH 20/28] Add changelog --- CHANGELOG.D/793.feature | 1 + 1 file changed, 1 insertion(+) create mode 100644 CHANGELOG.D/793.feature diff --git a/CHANGELOG.D/793.feature b/CHANGELOG.D/793.feature new file mode 100644 index 000000000..2a2449cc3 --- /dev/null +++ b/CHANGELOG.D/793.feature @@ -0,0 +1 @@ +Implement `neuro config login-headless` command. \ No newline at end of file From bd60351e514711f7869fd9262317cc08cc255e28 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 20 May 2019 13:45:12 +0300 Subject: [PATCH 21/28] Properly handle AuthCode cancellation --- neuromation/api/login.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/neuromation/api/login.py b/neuromation/api/login.py index f5498f4fc..3467a07f5 100644 --- a/neuromation/api/login.py +++ b/neuromation/api/login.py @@ -72,7 +72,8 @@ def challenge_method(self) -> str: return self._challenge_method def set_value(self, value: str) -> None: - self._future.set_result(value) + if not self._future.cancelled(): + self._future.set_result(value) @property def callback_url(self) -> URL: @@ -84,7 +85,8 @@ def callback_url(self, value: URL) -> None: self._callback_url = value def set_exception(self, exc: Exception) -> None: - self._future.set_exception(exc) + if not self._future.cancelled(): + self._future.set_exception(exc) def cancel(self) -> None: self._future.cancel() @@ -140,9 +142,14 @@ def __init__( self._callback = callback async def _request(self, url: URL, code: AuthCode) -> None: - auth_code = await self._callback(url) - assert auth_code - code.set_value(auth_code) + try: + auth_code = await self._callback(url) + assert auth_code + code.set_value(auth_code) + except asyncio.CancelledError: + raise + except Exception as exc: + code.set_exception(exc) class DummyAuthCodeCallbackClient(AuthCodeCallbackClient): From a600d6e4dcad142062697019fc4d9ee1c2b2fecf Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 20 May 2019 14:35:11 +0300 Subject: [PATCH 22/28] Work on --- tests/api/test_config_factory.py | 42 +++++++++++++++++++++++++++++--- tests/api/test_login.py | 12 ++++----- tests/api/test_login_utils.py | 6 ++--- 3 files changed, 47 insertions(+), 13 deletions(-) diff --git a/tests/api/test_config_factory.py b/tests/api/test_config_factory.py index fd2bc67c7..eec5a036f 100644 --- a/tests/api/test_config_factory.py +++ b/tests/api/test_config_factory.py @@ -1,6 +1,7 @@ import sys from pathlib import Path from typing import Any, List, Optional, Set +from unittest import mock from uuid import uuid4 as uuid import aiohttp @@ -57,7 +58,7 @@ async def _config_handler(request: web.Request) -> web.Response: "token_url": "https://test-neuromation.auth0.com/oauth/token", "client_id": "banana", "audience": "https://test.dev.neuromation.io", - "headless_callback_url": "https://https://dev.neu.ro/oauth/show-code", + "headless_callback_url": "https://dev.neu.ro/oauth/show-code", "callback_urls": [ "http://127.0.0.2:54540", "http://127.0.0.2:54541", @@ -110,7 +111,7 @@ def test_check_initialized(self) -> None: token_url=URL("http://token"), client_id="client-id", audience="everyone", - headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), + headless_callback_url=URL("https://dev.neu.ro/oauth/show-code"), ) assert auth_config_good.is_initialized() @@ -139,7 +140,7 @@ def test_check_initialized_bad_auth_config(self) -> None: token_url=URL("http://token"), client_id="client-id", audience="everyone", - headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), + headless_callback_url=URL("https://dev.neu.ro/oauth/show-code"), ) assert not auth_config_bad.is_initialized() @@ -169,7 +170,7 @@ def test_check_initialized_bad_cluster_config(self) -> None: token_url=URL("http://token"), client_id="client-id", audience="everyone", - headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), + headless_callback_url=URL("https://dev.neu.ro/oauth/show-code"), ) assert auth_config_bad.is_initialized() @@ -303,6 +304,39 @@ async def test_normal_login(self, tmp_home: Path, mock_for_login: URL) -> None: assert saved_config.cluster_config.is_initialized() +class TestHeadlessLogin: + async def test_login_headless_already_logged(self, config_file: Path) -> None: + async def callback(url: URL) -> str: + return "" + + with pytest.raises(ConfigError, match=r"already exists"): + await Factory().login_headless(callback) + + async def test_normal_login(self, tmp_home: Path, mock_for_login: URL) -> None: + async def callback(url: URL) -> str: + assert url.with_query(None) == URL( + "https://test-neuromation.auth0.com/authorize" + ) + + assert dict(url.query) == dict( + response_type="code", + code_challenge=mock.ANY, + code_challenge_method="S256", + client_id="banana", + redirect_uri="https://dev.neu.ro/oauth/show-code", + scope="offline_access", + audience="https://test.dev.neuromation.io", + ) + return "test_code" + + await Factory().login_headless(callback=callback, url=mock_for_login) + nmrc_path = tmp_home / ".nmrc" + assert Path(nmrc_path).exists(), "Config file not written after login " + saved_config = Factory(nmrc_path)._read() + assert saved_config.auth_config.is_initialized() + assert saved_config.cluster_config.is_initialized() + + class TestLogout: async def test_logout(self, config_file: Path) -> None: await Factory().logout() diff --git a/tests/api/test_login.py b/tests/api/test_login.py index ef1292584..bc2573f02 100644 --- a/tests/api/test_login.py +++ b/tests/api/test_login.py @@ -391,7 +391,7 @@ def test_is_initialized__no_auth_url(self) -> None: token_url=URL("url"), client_id="client_id", audience="audience", - headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), + headless_callback_url=URL("https://dev.neu.ro/oauth/show-code"), callback_urls=(URL("url1"), URL("url2")), success_redirect_url=URL("url"), ) @@ -403,7 +403,7 @@ def test_is_initialized__no_token_url(self) -> None: token_url=URL(), client_id="client_id", audience="audience", - headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), + headless_callback_url=URL("https://dev.neu.ro/oauth/show-code"), callback_urls=(URL("url1"), URL("url2")), success_redirect_url=URL("url"), ) @@ -415,7 +415,7 @@ def test_is_initialized__no_client_id(self) -> None: token_url=URL("url"), client_id="", audience="audience", - headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), + headless_callback_url=URL("https://dev.neu.ro/oauth/show-code"), callback_urls=(URL("url1"), URL("url2")), success_redirect_url=URL("url"), ) @@ -427,7 +427,7 @@ def test_is_initialized__no_audience(self) -> None: token_url=URL("url"), client_id="client_id", audience="", - headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), + headless_callback_url=URL("https://dev.neu.ro/oauth/show-code"), callback_urls=(URL("url1"), URL("url2")), success_redirect_url=URL("url"), ) @@ -439,7 +439,7 @@ def test_is_initialized__no_callback_urls(self) -> None: token_url=URL("url"), client_id="client_id", audience="audience", - headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), + headless_callback_url=URL("https://dev.neu.ro/oauth/show-code"), callback_urls=[], success_redirect_url=URL("url"), ) @@ -451,7 +451,7 @@ def test_is_initialized__no_success_redirect_url(self) -> None: token_url=URL("url"), client_id="client_id", audience="audience", - headless_callback_url=URL("https://https://dev.neu.ro/oauth/show-code"), + headless_callback_url=URL("https://dev.neu.ro/oauth/show-code"), callback_urls=(URL("url1"), URL("url2")), success_redirect_url=None, ) diff --git a/tests/api/test_login_utils.py b/tests/api/test_login_utils.py index e3e25dcf6..7422d5a96 100644 --- a/tests/api/test_login_utils.py +++ b/tests/api/test_login_utils.py @@ -17,7 +17,7 @@ async def test_get_server_config(aiohttp_server: _TestServerFactory) -> None: token_url = "https://dev-neuromation.auth0.com/oauth/token" client_id = "this_is_client_id" audience = "https://platform.dev.neuromation.io" - headless_callback_url = "https://https://dev.neu.ro/oauth/show-code" + headless_callback_url = "https://dev.neu.ro/oauth/show-code" callback_urls = [ "http://127.0.0.1:54540", "http://127.0.0.1:54541", @@ -66,7 +66,7 @@ async def test_get_server_config_no_callback_urls( token_url = "https://dev-neuromation.auth0.com/oauth/token" client_id = "this_is_client_id" audience = "https://platform.dev.neuromation.io" - headless_callback_url = "https://https://dev.neu.ro/oauth/show-code" + headless_callback_url = "https://dev.neu.ro/oauth/show-code" success_redirect_url = "https://platform.neuromation.io" JSON = { "auth_url": auth_url, @@ -110,7 +110,7 @@ async def test_get_server_config_with_token(aiohttp_server: _TestServerFactory) token_url = "https://dev-neuromation.auth0.com/oauth/token" client_id = "this_is_client_id" audience = "https://platform.dev.neuromation.io" - headless_callback_url = "https://https://dev.neu.ro/oauth/show-code" + headless_callback_url = "https://dev.neu.ro/oauth/show-code" success_redirect_url = "https://platform.neuromation.io" JSON = { "registry_url": registry_url, From d579107033fff01e5f1a095139dde55f85421fc8 Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 27 May 2019 12:46:47 +0300 Subject: [PATCH 23/28] Fix test --- neuromation/api/config_factory.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/neuromation/api/config_factory.py b/neuromation/api/config_factory.py index af1415688..57e6ba51f 100644 --- a/neuromation/api/config_factory.py +++ b/neuromation/api/config_factory.py @@ -67,8 +67,6 @@ async def login( pypi=_PyPIVersion.create_uninitialized(), url=url, ) - async with Client._create(config, timeout=timeout) as client: - await client.jobs.list() # raises an exception if cannot login self._save(config) async def login_headless( From 29a44bd0be85b30dba7a2f352a7922cdb47b3ded Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 27 May 2019 13:08:17 +0300 Subject: [PATCH 24/28] Use shared TCPConnector in get_server_config --- neuromation/api/config_factory.py | 29 ++++++++++++++++++----------- neuromation/api/login.py | 8 ++++++-- tests/api/test_login_utils.py | 14 ++++++++++---- 3 files changed, 34 insertions(+), 17 deletions(-) diff --git a/neuromation/api/config_factory.py b/neuromation/api/config_factory.py index 57e6ba51f..019146ffc 100644 --- a/neuromation/api/config_factory.py +++ b/neuromation/api/config_factory.py @@ -55,11 +55,14 @@ async def login( ) -> None: if self._path.exists(): raise ConfigError(f"Config file {self._path} already exists. Please logout") - config_unauthorized = await get_server_config(url) - negotiator = AuthNegotiator(config_unauthorized.auth_config, timeout) - auth_token = await negotiator.refresh_token() + async with aiohttp.TCPConnector() as connector: + config_unauthorized = await get_server_config(connector, url) + negotiator = AuthNegotiator(config_unauthorized.auth_config, timeout) + auth_token = await negotiator.refresh_token() - config_authorized = await get_server_config(url, token=auth_token.token) + config_authorized = await get_server_config( + connector, url, token=auth_token.token + ) config = _Config( auth_config=config_authorized.auth_config, auth_token=auth_token, @@ -78,13 +81,16 @@ async def login_headless( ) -> None: if self._path.exists(): raise ConfigError(f"Config file {self._path} already exists. Please logout") - config_unauthorized = await get_server_config(url) - negotiator = HeadlessNegotiator( - config_unauthorized.auth_config, callback, timeout - ) - auth_token = await negotiator.refresh_token() + async with aiohttp.TCPConnector() as connector: + config_unauthorized = await get_server_config(connector, url) + negotiator = HeadlessNegotiator( + config_unauthorized.auth_config, callback, timeout + ) + auth_token = await negotiator.refresh_token() - config_authorized = await get_server_config(url, token=auth_token.token) + config_authorized = await get_server_config( + connector, url, token=auth_token.token + ) config = _Config( auth_config=config_authorized.auth_config, auth_token=auth_token, @@ -103,7 +109,8 @@ async def login_with_token( ) -> None: if self._path.exists(): raise ConfigError(f"Config file {self._path} already exists. Please logout") - server_config = await get_server_config(url, token=token) + async with aiohttp.TCPConnector() as connector: + server_config = await get_server_config(connector, url, token=token) config = _Config( auth_config=server_config.auth_config, auth_token=_AuthToken.create_non_expiring(token), diff --git a/neuromation/api/login.py b/neuromation/api/login.py index 31bedcafb..fa1ff7c47 100644 --- a/neuromation/api/login.py +++ b/neuromation/api/login.py @@ -514,8 +514,12 @@ class ConfigLoadException(Exception): pass -async def get_server_config(url: URL, token: Optional[str] = None) -> _ServerConfig: - async with aiohttp.ClientSession(timeout=DEFAULT_TIMEOUT) as client: +async def get_server_config( + connector: aiohttp.BaseConnector, url: URL, token: Optional[str] = None +) -> _ServerConfig: + async with aiohttp.ClientSession( + timeout=DEFAULT_TIMEOUT, connector=connector, connector_owner=False + ) as client: headers: Dict[str, str] = {} if token: headers["Authorization"] = f"Bearer {token}" diff --git a/tests/api/test_login_utils.py b/tests/api/test_login_utils.py index 7422d5a96..cc860611b 100644 --- a/tests/api/test_login_utils.py +++ b/tests/api/test_login_utils.py @@ -42,7 +42,8 @@ async def handler(request: web.Request) -> web.Response: app.router.add_get("/config", handler) srv = await aiohttp_server(app) - config = await get_server_config(srv.make_url("/")) + async with aiohttp.TCPConnector() as connector: + config = await get_server_config(connector, srv.make_url("/")) assert config == _ServerConfig( auth_config=_AuthConfig( auth_url=URL(auth_url), @@ -85,7 +86,8 @@ async def handler(request: web.Request) -> web.Response: app.router.add_get("/config", handler) srv = await aiohttp_server(app) - config = await get_server_config(srv.make_url("/")) + async with aiohttp.TCPConnector() as connector: + config = await get_server_config(connector, srv.make_url("/")) assert config == _ServerConfig( auth_config=_AuthConfig( auth_url=URL(auth_url), @@ -133,7 +135,10 @@ async def handler(request: web.Request) -> web.Response: app.router.add_get("/config", handler) srv = await aiohttp_server(app) - config = await get_server_config(srv.make_url("/"), token="bananatoken") + async with aiohttp.TCPConnector() as connector: + config = await get_server_config( + connector, srv.make_url("/"), token="bananatoken" + ) assert config == _ServerConfig( auth_config=_AuthConfig( auth_url=URL(auth_url), @@ -161,4 +166,5 @@ async def handler(request: web.Request) -> web.Response: srv = await aiohttp_server(app) with pytest.raises(RuntimeError, match="Unable to get server configuration: 500"): - await get_server_config(srv.make_url("/")) + async with aiohttp.TCPConnector() as connector: + await get_server_config(connector, srv.make_url("/")) From 9e8c481811fa5482ade99d3a3bd776d6dbe8c01b Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 27 May 2019 15:09:17 +0300 Subject: [PATCH 25/28] Use explicit connector --- neuromation/api/__init__.py | 56 ++---------------------- neuromation/api/client.py | 14 +++--- neuromation/api/config_factory.py | 43 ++++++++++++++---- neuromation/api/core.py | 4 +- neuromation/api/login.py | 32 +++++++++++--- neuromation/api/registry.py | 2 +- neuromation/api/utils.py | 60 +++++++++++++++++++++++++- tests/api/test_config_factory.py | 9 ++-- tests/api/test_login.py | 72 +++++++++++++++++++++++++------ tests/conftest.py | 4 +- 10 files changed, 199 insertions(+), 97 deletions(-) diff --git a/neuromation/api/__init__.py b/neuromation/api/__init__.py index 898fd6fde..1f785e35f 100644 --- a/neuromation/api/__init__.py +++ b/neuromation/api/__init__.py @@ -1,7 +1,5 @@ -import sys from pathlib import Path -from types import TracebackType -from typing import Any, Awaitable, Callable, Coroutine, Generator, Optional, Type +from typing import Awaitable, Callable, Optional import aiohttp from yarl import URL @@ -40,26 +38,7 @@ from .parsing_utils import ImageNameParser from .storage import FileStatus, FileStatusType from .users import Action, Permission, SharedPermission - - -if sys.version_info >= (3, 7): - from typing import AsyncContextManager -else: - from typing import Generic, TypeVar - - _T = TypeVar("_T") - - class AsyncContextManager(Generic[_T]): - async def __aenter__(self) -> _T: - pass # pragma: no cover - - async def __aexit__( - self, - exc_type: Optional[Type[BaseException]], - exc: Optional[BaseException], - tb: Optional[TracebackType], - ) -> Optional[bool]: - pass # pragma: no cover +from .utils import _ContextManager __all__ = ( @@ -103,37 +82,10 @@ async def __aexit__( ) -class _ContextManager(Awaitable[Client], AsyncContextManager[Client]): - - __slots__ = ("_coro", "_client") - - def __init__(self, coro: Coroutine[Any, Any, Client]) -> None: - self._coro = coro - self._client: Optional[Client] = None - - def __await__(self) -> Generator[Any, None, Client]: - return self._coro.__await__() - - async def __aenter__(self) -> Client: - self._client = await self._coro - assert self._client is not None - return self._client - - async def __aexit__( - self, - exc_type: Optional[Type[BaseException]], - exc: Optional[BaseException], - tb: Optional[TracebackType], - ) -> Optional[bool]: - assert self._client is not None - await self._client.close() - return None - - def get( *, path: Optional[Path] = None, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT -) -> _ContextManager: - return _ContextManager(_get(path, timeout)) +) -> _ContextManager[Client]: + return _ContextManager[Client](_get(path, timeout)) async def _get(path: Optional[Path], timeout: aiohttp.ClientTimeout) -> Client: diff --git a/neuromation/api/client.py b/neuromation/api/client.py index 022d3268a..be8d20056 100644 --- a/neuromation/api/client.py +++ b/neuromation/api/client.py @@ -1,9 +1,7 @@ -import ssl from types import TracebackType from typing import Optional, Type import aiohttp -import certifi from .config import _Config from .core import DEFAULT_TIMEOUT, _Core @@ -16,15 +14,17 @@ class Client(metaclass=NoPublicConstructor): def __init__( - self, config: _Config, *, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT + self, + connector: aiohttp.BaseConnector, + config: _Config, + *, + timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT ) -> None: config.check_initialized() self._config = config - self._ssl_context = ssl.SSLContext() - self._ssl_context.load_verify_locations(capath=certifi.where()) - self._connector = aiohttp.TCPConnector(ssl=self._ssl_context) + self._connector = connector self._core = _Core( - self._connector, self._config.url, self._config.auth_token.token, timeout + connector, self._config.url, self._config.auth_token.token, timeout ) self._jobs = Jobs._create(self._core, self._config) self._storage = Storage._create(self._core, self._config) diff --git a/neuromation/api/config_factory.py b/neuromation/api/config_factory.py index 019146ffc..c954b08a1 100644 --- a/neuromation/api/config_factory.py +++ b/neuromation/api/config_factory.py @@ -1,10 +1,13 @@ +import asyncio import os +import ssl import sys from dataclasses import replace from pathlib import Path from typing import Any, Awaitable, Callable, Dict, Optional import aiohttp +import certifi import yaml from yarl import URL @@ -20,6 +23,7 @@ get_server_config, refresh_token, ) +from .utils import _ContextManager WIN32 = sys.platform == "win32" @@ -28,6 +32,17 @@ DEFAULT_API_URL = URL("https://staging.neu.ro/api/v1") +def _make_connector() -> _ContextManager[aiohttp.TCPConnector]: + return _ContextManager[aiohttp.TCPConnector](__make_connector()) + + +async def __make_connector() -> aiohttp.TCPConnector: + ssl_context = ssl.SSLContext() + ssl_context.load_verify_locations(capath=certifi.where()) + connector = aiohttp.TCPConnector(ssl=ssl_context) + return connector + + class ConfigError(RuntimeError): pass @@ -40,12 +55,22 @@ def __init__(self, path: Optional[Path] = None) -> None: async def get(self, *, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT) -> Client: config = self._read() - new_token = await refresh_token(config.auth_config, config.auth_token, timeout) + connector = await _make_connector() + try: + new_token = await refresh_token( + connector, config.auth_config, config.auth_token, timeout + ) + except asyncio.CancelledError: + await connector.close() + raise + except Exception: + await connector.close() + raise if new_token != config.auth_token: new_config = replace(config, auth_token=new_token) self._save(new_config) - return Client._create(new_config, timeout=timeout) - return Client._create(config, timeout=timeout) + return Client._create(connector, new_config, timeout=timeout) + return Client._create(connector, config, timeout=timeout) async def login( self, @@ -55,9 +80,11 @@ async def login( ) -> None: if self._path.exists(): raise ConfigError(f"Config file {self._path} already exists. Please logout") - async with aiohttp.TCPConnector() as connector: + async with _make_connector() as connector: config_unauthorized = await get_server_config(connector, url) - negotiator = AuthNegotiator(config_unauthorized.auth_config, timeout) + negotiator = AuthNegotiator( + connector, config_unauthorized.auth_config, timeout + ) auth_token = await negotiator.refresh_token() config_authorized = await get_server_config( @@ -81,10 +108,10 @@ async def login_headless( ) -> None: if self._path.exists(): raise ConfigError(f"Config file {self._path} already exists. Please logout") - async with aiohttp.TCPConnector() as connector: + async with _make_connector() as connector: config_unauthorized = await get_server_config(connector, url) negotiator = HeadlessNegotiator( - config_unauthorized.auth_config, callback, timeout + connector, config_unauthorized.auth_config, callback, timeout ) auth_token = await negotiator.refresh_token() @@ -109,7 +136,7 @@ async def login_with_token( ) -> None: if self._path.exists(): raise ConfigError(f"Config file {self._path} already exists. Please logout") - async with aiohttp.TCPConnector() as connector: + async with _make_connector() as connector: server_config = await get_server_config(connector, url, token=token) config = _Config( auth_config=server_config.auth_config, diff --git a/neuromation/api/core.py b/neuromation/api/core.py index 951426a52..f081fbdc5 100644 --- a/neuromation/api/core.py +++ b/neuromation/api/core.py @@ -45,7 +45,7 @@ class _Core: def __init__( self, - connector: aiohttp.TCPConnector, + connector: aiohttp.BaseConnector, base_url: URL, token: str, timeout: aiohttp.ClientTimeout, @@ -69,7 +69,7 @@ def __init__( } @property - def connector(self) -> aiohttp.TCPConnector: + def connector(self) -> aiohttp.BaseConnector: return self._connector @property diff --git a/neuromation/api/login.py b/neuromation/api/login.py index fa1ff7c47..73f7548e8 100644 --- a/neuromation/api/login.py +++ b/neuromation/api/login.py @@ -283,12 +283,18 @@ def create_non_expiring(cls, token: str) -> "_AuthToken": class AuthTokenClient: def __init__( - self, url: URL, client_id: str, timeout: aiohttp.ClientTimeout + self, + connector: aiohttp.BaseConnector, + url: URL, + client_id: str, + timeout: aiohttp.ClientTimeout, ) -> None: self._url = url self._client_id = client_id - self._client = ClientSession(timeout=timeout) + self._client = ClientSession( + connector=connector, connector_owner=False, timeout=timeout + ) async def close(self) -> None: await self._client.close() @@ -418,10 +424,13 @@ def create( async def refresh_token( - config: _AuthConfig, token: _AuthToken, timeout: aiohttp.ClientTimeout + connector: aiohttp.BaseConnector, + config: _AuthConfig, + token: _AuthToken, + timeout: aiohttp.ClientTimeout, ) -> _AuthToken: async with AuthTokenClient( - url=config.token_url, client_id=config.client_id, timeout=timeout + connector, url=config.token_url, client_id=config.client_id, timeout=timeout ) as token_client: if token.is_expired: return await token_client.refresh(token) @@ -429,9 +438,15 @@ async def refresh_token( class BaseNegotiator(abc.ABC): - def __init__(self, config: _AuthConfig, timeout: aiohttp.ClientTimeout) -> None: + def __init__( + self, + connector: aiohttp.BaseConnector, + config: _AuthConfig, + timeout: aiohttp.ClientTimeout, + ) -> None: self._config = config self._timeout = timeout + self._connector = connector @abc.abstractmethod async def get_code(self) -> AuthCode: @@ -439,6 +454,7 @@ async def get_code(self) -> AuthCode: async def refresh_token(self, token: Optional[_AuthToken] = None) -> _AuthToken: async with AuthTokenClient( + self._connector, url=self._config.token_url, client_id=self._config.client_id, timeout=self._timeout, @@ -456,13 +472,14 @@ async def refresh_token(self, token: Optional[_AuthToken] = None) -> _AuthToken: class AuthNegotiator(BaseNegotiator): def __init__( self, + connector: aiohttp.BaseConnector, config: _AuthConfig, timeout: aiohttp.ClientTimeout, code_callback_client_factory: Type[ AuthCodeCallbackClient ] = WebBrowserAuthCodeCallbackClient, ) -> None: - super().__init__(config, timeout) + super().__init__(connector, config, timeout) self._code_callback_client_factory = code_callback_client_factory async def get_code(self) -> AuthCode: @@ -484,11 +501,12 @@ async def get_code(self) -> AuthCode: class HeadlessNegotiator(BaseNegotiator): def __init__( self, + connector: aiohttp.BaseConnector, config: _AuthConfig, callback: Callable[[URL], Awaitable[str]], timeout: aiohttp.ClientTimeout, ) -> None: - super().__init__(config, timeout) + super().__init__(connector, config, timeout) self._callback = callback async def get_code(self) -> AuthCode: diff --git a/neuromation/api/registry.py b/neuromation/api/registry.py index 4ee4c7fa1..5cb5bec87 100644 --- a/neuromation/api/registry.py +++ b/neuromation/api/registry.py @@ -17,7 +17,7 @@ class _Registry(_Core): """ def __init__( - self, connector: aiohttp.TCPConnector, base_url: URL, token: str, username: str + self, connector: aiohttp.BaseConnector, base_url: URL, token: str, username: str ) -> None: self._username = username super().__init__(connector, base_url, token, TIMEOUT) diff --git a/neuromation/api/utils.py b/neuromation/api/utils.py index 35ef6ea11..f88c68e79 100644 --- a/neuromation/api/utils.py +++ b/neuromation/api/utils.py @@ -1,5 +1,15 @@ import sys -from typing import Any +from types import TracebackType +from typing import ( + Any, + Awaitable, + Coroutine, + Generator, + Generic, + Optional, + Type, + TypeVar, +) if sys.version_info >= (3, 7): # pragma: no cover @@ -7,6 +17,25 @@ else: from async_generator import asynccontextmanager # noqa +_T = TypeVar("_T") + + +if sys.version_info >= (3, 7): + from typing import AsyncContextManager +else: + + class AsyncContextManager(Generic[_T]): + async def __aenter__(self) -> _T: + pass # pragma: no cover + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc: Optional[BaseException], + tb: Optional[TracebackType], + ) -> Optional[bool]: + pass # pragma: no cover + class NoPublicConstructor(type): def __call__(self, *args: Any, **kwargs: Any) -> Any: @@ -15,3 +44,32 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: def _create(self, *args: Any, **kwargs: Any) -> Any: return super().__call__(*args, **kwargs) + + +class _ContextManager(Generic[_T], Awaitable[_T], AsyncContextManager[_T]): + + __slots__ = ("_coro", "_ret") + + def __init__(self, coro: Coroutine[Any, Any, _T]) -> None: + self._coro = coro + self._ret: Optional[_T] = None + + def __await__(self) -> Generator[Any, None, _T]: + return self._coro.__await__() + + async def __aenter__(self) -> _T: + self._ret = await self._coro + assert self._ret is not None + return self._ret + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc: Optional[BaseException], + tb: Optional[TracebackType], + ) -> Optional[bool]: + assert self._ret is not None + # ret supports async close() protocol + # Need to teach mypy about this facility + await self._ret.close() # type: ignore + return None diff --git a/tests/api/test_config_factory.py b/tests/api/test_config_factory.py index 143df22d7..a0dd7c7be 100644 --- a/tests/api/test_config_factory.py +++ b/tests/api/test_config_factory.py @@ -41,9 +41,7 @@ def config_file( @pytest.fixture async def mock_for_login(monkeypatch: Any, aiohttp_server: _TestServerFactory) -> URL: - async def _refresh_token_mock( - self: Any, token: Optional[_AuthToken] = None - ) -> _AuthToken: + async def _refresh_token_mock(*args: Any, **kwargs: Any) -> _AuthToken: return _AuthToken.create_non_expiring(str(uuid())) async def _config_handler(request: web.Request) -> web.Response: @@ -232,7 +230,10 @@ async def test_token_autorefreshing( new_token = str(uuid()) + "changed" * 10 # token must has other size async def _refresh_token_mock( - configf: _AuthConfig, token: _AuthToken, timeout: aiohttp.ClientTimeout + connector: aiohttp.BaseConnector, + config: _AuthConfig, + token: _AuthToken, + timeout: aiohttp.ClientTimeout, ) -> _AuthToken: return _AuthToken.create_non_expiring(new_token) diff --git a/tests/api/test_login.py b/tests/api/test_login.py index bc2573f02..dd85f923c 100644 --- a/tests/api/test_login.py +++ b/tests/api/test_login.py @@ -2,6 +2,7 @@ from typing import AsyncIterator, Optional from unittest import mock +import aiohttp import pytest from aiohttp import ClientSession from aiohttp.test_utils import unused_port @@ -36,6 +37,14 @@ from tests import _TestServerFactory +@pytest.fixture +async def connector( + loop: asyncio.AbstractEventLoop +) -> AsyncIterator[aiohttp.BaseConnector]: + async with aiohttp.TCPConnector() as connector: + yield connector + + class TestAuthCode: async def test_wait_timed_out(self) -> None: code = AuthCode() @@ -323,20 +332,33 @@ async def auth_config( class TestTokenClient: - async def test_request(self, auth_client_id: str, auth_config: _AuthConfig) -> None: + async def test_request( + self, + connector: aiohttp.BaseConnector, + auth_client_id: str, + auth_config: _AuthConfig, + ) -> None: code = AuthCode() code.set_value("test_code") code.callback_url = auth_config.callback_urls[0] async with AuthTokenClient( - auth_config.token_url, client_id=auth_client_id, timeout=DEFAULT_TIMEOUT + connector, + auth_config.token_url, + client_id=auth_client_id, + timeout=DEFAULT_TIMEOUT, ) as client: token = await client.request(code) assert token.token == "test_access_token" assert token.refresh_token == "test_refresh_token" assert not token.is_expired - async def test_refresh(self, auth_client_id: str, auth_config: _AuthConfig) -> None: + async def test_refresh( + self, + connector: aiohttp.BaseConnector, + auth_client_id: str, + auth_config: _AuthConfig, + ) -> None: token = _AuthToken.create( token="test_access_token", expires_in=1234, @@ -344,7 +366,10 @@ async def test_refresh(self, auth_client_id: str, auth_config: _AuthConfig) -> N ) async with AuthTokenClient( - auth_config.token_url, client_id=auth_client_id, timeout=DEFAULT_TIMEOUT + connector, + auth_config.token_url, + client_id=auth_client_id, + timeout=DEFAULT_TIMEOUT, ) as client: new_token = await client.refresh(token) assert new_token.token == "test_access_token_refreshed" @@ -352,7 +377,10 @@ async def test_refresh(self, auth_client_id: str, auth_config: _AuthConfig) -> N assert not token.is_expired async def test_forbidden( - self, aiohttp_server: _TestServerFactory, auth_config: _AuthConfig + self, + connector: aiohttp.BaseConnector, + aiohttp_server: _TestServerFactory, + auth_config: _AuthConfig, ) -> None: code = AuthCode() code.callback_url = auth_config.callback_urls[0] @@ -370,7 +398,7 @@ async def handle_token(request: Request) -> Response: url = server.make_url("/oauth/token") async with AuthTokenClient( - url, client_id=client_id, timeout=DEFAULT_TIMEOUT + connector, url, client_id=client_id, timeout=DEFAULT_TIMEOUT ) as client: with pytest.raises(AuthException, match="failed to get an access token."): await client.request(code) @@ -518,8 +546,11 @@ def test_is_initialized__no_monitoring_url(self) -> None: class TestAuthNegotiator: - async def test_get_code(self, auth_config: _AuthConfig) -> None: + async def test_get_code( + self, connector: aiohttp.BaseConnector, auth_config: _AuthConfig + ) -> None: negotiator = AuthNegotiator( + connector, config=auth_config, timeout=DEFAULT_TIMEOUT, code_callback_client_factory=DummyAuthCodeCallbackClient, @@ -528,8 +559,11 @@ async def test_get_code(self, auth_config: _AuthConfig) -> None: assert await code.wait() == "test_code" assert code.callback_url == auth_config.callback_urls[0] - async def test_get_token(self, auth_config: _AuthConfig) -> None: + async def test_get_token( + self, connector: aiohttp.BaseConnector, auth_config: _AuthConfig + ) -> None: negotiator = AuthNegotiator( + connector, config=auth_config, timeout=DEFAULT_TIMEOUT, code_callback_client_factory=DummyAuthCodeCallbackClient, @@ -538,8 +572,11 @@ async def test_get_token(self, auth_config: _AuthConfig) -> None: assert token.token == "test_access_token" assert token.refresh_token == "test_refresh_token" - async def test_refresh_token_noop(self, auth_config: _AuthConfig) -> None: + async def test_refresh_token_noop( + self, connector: aiohttp.BaseConnector, auth_config: _AuthConfig + ) -> None: negotiator = AuthNegotiator( + connector, config=auth_config, timeout=DEFAULT_TIMEOUT, code_callback_client_factory=DummyAuthCodeCallbackClient, @@ -553,8 +590,11 @@ async def test_refresh_token_noop(self, auth_config: _AuthConfig) -> None: assert token.token == "test_access_token" assert token.refresh_token == "test_refresh_token" - async def test_refresh_token(self, auth_config: _AuthConfig) -> None: + async def test_refresh_token( + self, connector: aiohttp.BaseConnector, auth_config: _AuthConfig + ) -> None: negotiator = AuthNegotiator( + connector, config=auth_config, timeout=DEFAULT_TIMEOUT, code_callback_client_factory=DummyAuthCodeCallbackClient, @@ -573,7 +613,9 @@ async def test_refresh_token(self, auth_config: _AuthConfig) -> None: class TestHeadlessNegotiator: - async def test_get_code(self, auth_config: _AuthConfig) -> None: + async def test_get_code( + self, connector: aiohttp.BaseConnector, auth_config: _AuthConfig + ) -> None: async def callback(url: URL) -> str: assert url.with_query(None) == auth_config.auth_url @@ -589,17 +631,19 @@ async def callback(url: URL) -> str: return "test_code" negotiator = HeadlessNegotiator( - config=auth_config, callback=callback, timeout=DEFAULT_TIMEOUT + connector, config=auth_config, callback=callback, timeout=DEFAULT_TIMEOUT ) code = await negotiator.get_code() assert await code.wait() == "test_code" - async def test_get_code_raises(self, auth_config: _AuthConfig) -> None: + async def test_get_code_raises( + self, connector: aiohttp.BaseConnector, auth_config: _AuthConfig + ) -> None: async def callback(url: URL) -> str: raise RuntimeError("callback error") negotiator = HeadlessNegotiator( - config=auth_config, callback=callback, timeout=DEFAULT_TIMEOUT + connector, config=auth_config, callback=callback, timeout=DEFAULT_TIMEOUT ) with pytest.raises(RuntimeError, match="callback error"): await negotiator.get_code() diff --git a/tests/conftest.py b/tests/conftest.py index e07ffc96f..665c7931f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,6 @@ from typing import Callable +import aiohttp import pytest from jose import jwt from yarl import URL @@ -58,6 +59,7 @@ def go(url_str: str, registry_url: str = "https://registry-dev.neu.ro") -> Clien url=URL(url), cluster_config=cluster_config, ) - return Client._create(config) + connector = aiohttp.TCPConnector() + return Client._create(connector, config) return go From 4f05055d39acf45797c734d92a2de4bfe16265be Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Mon, 27 May 2019 17:06:24 +0300 Subject: [PATCH 26/28] Work on --- tests/api/test_config_factory.py | 64 ++++++++++++++++++++++++-------- 1 file changed, 48 insertions(+), 16 deletions(-) diff --git a/tests/api/test_config_factory.py b/tests/api/test_config_factory.py index a0dd7c7be..fcb1bffff 100644 --- a/tests/api/test_config_factory.py +++ b/tests/api/test_config_factory.py @@ -8,6 +8,7 @@ import pytest import yaml from aiohttp import web +from aiohttp.test_utils import TestServer as _TestServer from yarl import URL import neuromation.api.config_factory @@ -40,17 +41,19 @@ def config_file( @pytest.fixture -async def mock_for_login(monkeypatch: Any, aiohttp_server: _TestServerFactory) -> URL: +async def mock_for_login( + monkeypatch: Any, aiohttp_server: _TestServerFactory +) -> _TestServer: async def _refresh_token_mock(*args: Any, **kwargs: Any) -> _AuthToken: return _AuthToken.create_non_expiring(str(uuid())) - async def _config_handler(request: web.Request) -> web.Response: + async def config_handler(request: web.Request) -> web.Response: config_json = { - "auth_url": "https://test-neuromation.auth0.com/authorize", - "token_url": "https://test-neuromation.auth0.com/oauth/token", + "auth_url": str(srv.make_url("/authorize")), + "token_url": str(srv.make_url("/oauth/token")), "client_id": "banana", "audience": "https://test.dev.neuromation.io", - "headless_callback_url": "https://dev.neu.ro/oauth/show-code", + "headless_callback_url": str(srv.make_url("/oauth/show-code")), "callback_urls": [ "http://127.0.0.2:54540", "http://127.0.0.2:54541", @@ -73,13 +76,28 @@ async def _config_handler(request: web.Request) -> web.Response: ) return web.json_response(config_json) + async def show_code(request: web.Request) -> web.Response: + breakpoint() + return web.json_response({}) + + async def authorize(request: web.Request) -> web.Response: + breakpoint() + return web.json_response({}) + + async def token(request: web.Request) -> web.Response: + breakpoint() + return web.json_response({}) + app = web.Application() - app.router.add_get("/config", _config_handler) + app.router.add_get("/config", config_handler) + app.router.add_get("/oauth/show-code", show_code) + app.router.add_get("/authorize", authorize) + app.router.add_get("/oauth/token", token) srv = await aiohttp_server(app) monkeypatch.setattr(AuthNegotiator, "refresh_token", _refresh_token_mock) - return srv.make_url("/") + return srv def _create_config( @@ -278,8 +296,10 @@ async def test_login_already_logged(self, config_file: Path) -> None: with pytest.raises(ConfigError, match=r"already exists"): await Factory().login() - async def test_normal_login(self, tmp_home: Path, mock_for_login: URL) -> None: - await Factory().login(url=mock_for_login) + async def test_normal_login( + self, tmp_home: Path, mock_for_login: _TestServer + ) -> None: + await Factory().login(url=mock_for_login.make_url("/")) nmrc_path = tmp_home / ".nmrc" assert Path(nmrc_path).exists(), "Config file not written after login " saved_config = Factory(nmrc_path)._read() @@ -292,17 +312,25 @@ async def test_login_with_token_already_logged(self, config_file: Path) -> None: with pytest.raises(ConfigError, match=r"already exists"): await Factory().login_with_token(token="tokenstr") - async def test_normal_login(self, tmp_home: Path, mock_for_login: URL) -> None: - await Factory().login_with_token(token="tokenstr", url=mock_for_login) + async def test_normal_login( + self, tmp_home: Path, mock_for_login: _TestServer + ) -> None: + await Factory().login_with_token( + token="tokenstr", url=mock_for_login.make_url("/") + ) nmrc_path = tmp_home / ".nmrc" assert Path(nmrc_path).exists(), "Config file not written after login " saved_config = Factory(nmrc_path)._read() assert saved_config.auth_config.is_initialized() assert saved_config.cluster_config.is_initialized() - async def test_incorrect_token(self, tmp_home: Path, mock_for_login: URL) -> None: + async def test_incorrect_token( + self, tmp_home: Path, mock_for_login: _TestServer + ) -> None: with pytest.raises(AuthException): - await Factory().login_with_token(token="incorrect", url=mock_for_login) + await Factory().login_with_token( + token="incorrect", url=mock_for_login.make_url("/") + ) nmrc_path = tmp_home / ".nmrc" assert not Path(nmrc_path).exists(), "Config file not written after login " @@ -315,7 +343,9 @@ async def callback(url: URL) -> str: with pytest.raises(ConfigError, match=r"already exists"): await Factory().login_headless(callback) - async def test_normal_login(self, tmp_home: Path, mock_for_login: URL) -> None: + async def test_normal_login( + self, tmp_home: Path, mock_for_login: _TestServer + ) -> None: async def callback(url: URL) -> str: assert url.with_query(None) == URL( "https://test-neuromation.auth0.com/authorize" @@ -326,13 +356,15 @@ async def callback(url: URL) -> str: code_challenge=mock.ANY, code_challenge_method="S256", client_id="banana", - redirect_uri="https://dev.neu.ro/oauth/show-code", + redirect_uri=mock_for_login.make_url("/oauth/show-code"), scope="offline_access", audience="https://test.dev.neuromation.io", ) return "test_code" - await Factory().login_headless(callback=callback, url=mock_for_login) + await Factory().login_headless( + callback=callback, url=mock_for_login.make_url("/") + ) nmrc_path = tmp_home / ".nmrc" assert Path(nmrc_path).exists(), "Config file not written after login " saved_config = Factory(nmrc_path)._read() From 944440877f991ec4ffbaf4e3a318a088669dfaec Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Tue, 28 May 2019 14:02:41 +0300 Subject: [PATCH 27/28] Work on --- neuromation/api/__init__.py | 7 +++--- neuromation/api/config_factory.py | 7 +++--- neuromation/api/login.py | 41 ++++++++++++++++--------------- neuromation/cli/config.py | 24 ++++++++++++------ tests/api/test_config_factory.py | 18 ++++++++------ tests/api/test_login.py | 27 +++++++++++++------- 6 files changed, 75 insertions(+), 49 deletions(-) diff --git a/neuromation/api/__init__.py b/neuromation/api/__init__.py index 1f785e35f..2359e0857 100644 --- a/neuromation/api/__init__.py +++ b/neuromation/api/__init__.py @@ -93,12 +93,13 @@ async def _get(path: Optional[Path], timeout: aiohttp.ClientTimeout) -> Client: async def login( + show_browser_cb: Callable[[URL], Awaitable[None]], *, url: URL = DEFAULT_API_URL, path: Optional[Path] = None, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT ) -> None: - await Factory(path).login(url=url, timeout=timeout) + await Factory(path).login(show_browser_cb, url=url, timeout=timeout) async def login_with_token( @@ -112,13 +113,13 @@ async def login_with_token( async def login_headless( - callback: Callable[[URL], Awaitable[str]], + get_auth_code_cb: Callable[[URL], Awaitable[str]], *, url: URL = DEFAULT_API_URL, path: Optional[Path] = None, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT ) -> None: - await Factory(path).login_headless(callback, url=url, timeout=timeout) + await Factory(path).login_headless(get_auth_code_cb, url=url, timeout=timeout) async def logout(*, path: Optional[Path] = None) -> None: diff --git a/neuromation/api/config_factory.py b/neuromation/api/config_factory.py index c954b08a1..658252c3a 100644 --- a/neuromation/api/config_factory.py +++ b/neuromation/api/config_factory.py @@ -74,6 +74,7 @@ async def get(self, *, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT) -> Clie async def login( self, + show_browser_cb: Callable[[URL], Awaitable[None]], *, url: URL = DEFAULT_API_URL, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT, @@ -83,7 +84,7 @@ async def login( async with _make_connector() as connector: config_unauthorized = await get_server_config(connector, url) negotiator = AuthNegotiator( - connector, config_unauthorized.auth_config, timeout + connector, config_unauthorized.auth_config, show_browser_cb, timeout ) auth_token = await negotiator.refresh_token() @@ -101,7 +102,7 @@ async def login( async def login_headless( self, - callback: Callable[[URL], Awaitable[str]], + get_auth_code_cb: Callable[[URL], Awaitable[str]], *, url: URL = DEFAULT_API_URL, timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT, @@ -111,7 +112,7 @@ async def login_headless( async with _make_connector() as connector: config_unauthorized = await get_server_config(connector, url) negotiator = HeadlessNegotiator( - connector, config_unauthorized.auth_config, callback, timeout + connector, config_unauthorized.auth_config, get_auth_code_cb, timeout ) auth_token = await negotiator.refresh_token() diff --git a/neuromation/api/login.py b/neuromation/api/login.py index 73f7548e8..b02d64b8d 100644 --- a/neuromation/api/login.py +++ b/neuromation/api/login.py @@ -5,7 +5,6 @@ import hashlib import secrets import time -import webbrowser from dataclasses import dataclass, field from typing import ( Any, @@ -125,9 +124,18 @@ async def _request(self, url: URL, code: AuthCode) -> None: class WebBrowserAuthCodeCallbackClient(AuthCodeCallbackClient): + def __init__( + self, + url: URL, + client_id: str, + audience: str, + show_browser_cb: Callable[[URL], Awaitable[None]], + ) -> None: + super().__init__(url=url, client_id=client_id, audience=audience) + self._show_browser_cb = show_browser_cb + async def _request(self, url: URL, code: AuthCode) -> None: - loop = asyncio.get_event_loop() - await loop.run_in_executor(None, webbrowser.open_new, str(url)) + await self._show_browser_cb(url) class HeadlessAuthCodeCallbackClient(AuthCodeCallbackClient): @@ -136,14 +144,14 @@ def __init__( url: URL, client_id: str, audience: str, - callback: Callable[[URL], Awaitable[str]], + get_auth_code_cb: Callable[[URL], Awaitable[str]], ) -> None: super().__init__(url=url, client_id=client_id, audience=audience) - self._callback = callback + self._get_auth_code_cb = get_auth_code_cb async def _request(self, url: URL, code: AuthCode) -> None: try: - auth_code = await self._callback(url) + auth_code = await self._get_auth_code_cb(url) assert auth_code code.set_value(auth_code) except asyncio.CancelledError: @@ -152,12 +160,6 @@ async def _request(self, url: URL, code: AuthCode) -> None: code.set_exception(exc) -class DummyAuthCodeCallbackClient(AuthCodeCallbackClient): - async def _request(self, url: URL, code: AuthCode) -> None: - async with ClientSession() as client: - await client.get(url, allow_redirects=True) - - class AuthCodeCallbackHandler: def __init__(self, code: AuthCode, redirect_url: Optional[URL] = None) -> None: self._code = code @@ -474,13 +476,11 @@ def __init__( self, connector: aiohttp.BaseConnector, config: _AuthConfig, + show_browser_cb: Callable[[URL], Awaitable[None]], timeout: aiohttp.ClientTimeout, - code_callback_client_factory: Type[ - AuthCodeCallbackClient - ] = WebBrowserAuthCodeCallbackClient, ) -> None: super().__init__(connector, config, timeout) - self._code_callback_client_factory = code_callback_client_factory + self._show_browser_cb = show_browser_cb async def get_code(self) -> AuthCode: code = AuthCode() @@ -490,10 +490,11 @@ async def get_code(self) -> AuthCode: app, host=self._config.callback_host, ports=self._config.callback_ports ) as url: code.callback_url = url - code_callback_client = self._code_callback_client_factory( + code_callback_client = WebBrowserAuthCodeCallbackClient( url=self._config.auth_url, client_id=self._config.client_id, audience=self._config.audience, + show_browser_cb=self._show_browser_cb, ) return await code_callback_client.request(code) @@ -503,11 +504,11 @@ def __init__( self, connector: aiohttp.BaseConnector, config: _AuthConfig, - callback: Callable[[URL], Awaitable[str]], + get_auth_code_cb: Callable[[URL], Awaitable[str]], timeout: aiohttp.ClientTimeout, ) -> None: super().__init__(connector, config, timeout) - self._callback = callback + self._get_auth_code_cb = get_auth_code_cb async def get_code(self) -> AuthCode: code = AuthCode() @@ -517,7 +518,7 @@ async def get_code(self) -> AuthCode: url=self._config.auth_url, client_id=self._config.client_id, audience=self._config.audience, - callback=self._callback, + get_auth_code_cb=self._get_auth_code_cb, ) return await code_callback_client.request(code) diff --git a/neuromation/cli/config.py b/neuromation/cli/config.py index dec804f2f..b9b8d7e87 100644 --- a/neuromation/cli/config.py +++ b/neuromation/cli/config.py @@ -1,6 +1,8 @@ +import asyncio import json import os import sys +import webbrowser from pathlib import Path from typing import Any, Dict @@ -54,12 +56,21 @@ async def login(root: Root, url: URL) -> None: URL is a platform entrypoint URL. """ + + async def show_browser(url: URL) -> None: + loop = asyncio.get_event_loop() + await loop.run_in_executor(None, webbrowser.open_new, str(url)) + try: - await api_login(url=url, path=root.config_path, timeout=root.timeout) + await api_login( + show_browser, url=url, path=root.config_path, timeout=root.timeout + ) except ConfigError: await api_logout(path=root.config_path) click.echo("You were successfully logged out.") - await api_login(url=url, path=root.config_path, timeout=root.timeout) + await api_login( + show_browser, url=url, path=root.config_path, timeout=root.timeout + ) click.echo(f"Logged into {url}") @@ -114,15 +125,14 @@ async def login_callback(url: URL) -> str: try: await api_login_headless( - url=url, - callback=login_callback, - path=root.config_path, - timeout=root.timeout, + login_callback, url=url, path=root.config_path, timeout=root.timeout ) except ConfigError: await api_logout(path=root.config_path) click.echo("You were successfully logged out.") - await api_login(url=url, path=root.config_path, timeout=root.timeout) + await api_login_headless( + login_callback, url=url, path=root.config_path, timeout=root.timeout + ) click.echo(f"Logged into {url}") diff --git a/tests/api/test_config_factory.py b/tests/api/test_config_factory.py index fcb1bffff..3c776bc90 100644 --- a/tests/api/test_config_factory.py +++ b/tests/api/test_config_factory.py @@ -1,6 +1,6 @@ import sys from pathlib import Path -from typing import Any, Optional +from typing import Any from unittest import mock from uuid import uuid4 as uuid @@ -292,14 +292,18 @@ async def test_mailformed_config(self, config_file: Path) -> None: class TestLogin: + async def show_dummy_browser(self, url: URL) -> None: + async with aiohttp.ClientSession() as client: + await client.get(url, allow_redirects=True) + async def test_login_already_logged(self, config_file: Path) -> None: with pytest.raises(ConfigError, match=r"already exists"): - await Factory().login() + await Factory().login(self.show_dummy_browser) async def test_normal_login( self, tmp_home: Path, mock_for_login: _TestServer ) -> None: - await Factory().login(url=mock_for_login.make_url("/")) + await Factory().login(self.show_dummy_browser, url=mock_for_login.make_url("/")) nmrc_path = tmp_home / ".nmrc" assert Path(nmrc_path).exists(), "Config file not written after login " saved_config = Factory(nmrc_path)._read() @@ -337,16 +341,16 @@ async def test_incorrect_token( class TestHeadlessLogin: async def test_login_headless_already_logged(self, config_file: Path) -> None: - async def callback(url: URL) -> str: + async def get_auth_code_cb(url: URL) -> str: return "" with pytest.raises(ConfigError, match=r"already exists"): - await Factory().login_headless(callback) + await Factory().login_headless(get_auth_code_cb) async def test_normal_login( self, tmp_home: Path, mock_for_login: _TestServer ) -> None: - async def callback(url: URL) -> str: + async def get_auth_code_cb(url: URL) -> str: assert url.with_query(None) == URL( "https://test-neuromation.auth0.com/authorize" ) @@ -363,7 +367,7 @@ async def callback(url: URL) -> str: return "test_code" await Factory().login_headless( - callback=callback, url=mock_for_login.make_url("/") + get_auth_code_cb, url=mock_for_login.make_url("/") ) nmrc_path = tmp_home / ".nmrc" assert Path(nmrc_path).exists(), "Config file not written after login " diff --git a/tests/api/test_login.py b/tests/api/test_login.py index dd85f923c..e0445a18a 100644 --- a/tests/api/test_login.py +++ b/tests/api/test_login.py @@ -25,7 +25,6 @@ AuthException, AuthNegotiator, AuthTokenClient, - DummyAuthCodeCallbackClient, HeadlessNegotiator, _AuthConfig, _AuthToken, @@ -546,14 +545,18 @@ def test_is_initialized__no_monitoring_url(self) -> None: class TestAuthNegotiator: + async def show_dummy_browser(self, url: URL) -> None: + async with ClientSession() as client: + await client.get(url, allow_redirects=True) + async def test_get_code( self, connector: aiohttp.BaseConnector, auth_config: _AuthConfig ) -> None: negotiator = AuthNegotiator( connector, config=auth_config, + show_browser_cb=self.show_dummy_browser, timeout=DEFAULT_TIMEOUT, - code_callback_client_factory=DummyAuthCodeCallbackClient, ) code = await negotiator.get_code() assert await code.wait() == "test_code" @@ -565,8 +568,8 @@ async def test_get_token( negotiator = AuthNegotiator( connector, config=auth_config, + show_browser_cb=self.show_dummy_browser, timeout=DEFAULT_TIMEOUT, - code_callback_client_factory=DummyAuthCodeCallbackClient, ) token = await negotiator.refresh_token(token=None) assert token.token == "test_access_token" @@ -578,8 +581,8 @@ async def test_refresh_token_noop( negotiator = AuthNegotiator( connector, config=auth_config, + show_browser_cb=self.show_dummy_browser, timeout=DEFAULT_TIMEOUT, - code_callback_client_factory=DummyAuthCodeCallbackClient, ) token = await negotiator.refresh_token(token=None) assert token.token == "test_access_token" @@ -596,8 +599,8 @@ async def test_refresh_token( negotiator = AuthNegotiator( connector, config=auth_config, + show_browser_cb=self.show_dummy_browser, timeout=DEFAULT_TIMEOUT, - code_callback_client_factory=DummyAuthCodeCallbackClient, ) token = await negotiator.refresh_token(token=None) assert token.token == "test_access_token" @@ -616,7 +619,7 @@ class TestHeadlessNegotiator: async def test_get_code( self, connector: aiohttp.BaseConnector, auth_config: _AuthConfig ) -> None: - async def callback(url: URL) -> str: + async def get_auth_code_cb(url: URL) -> str: assert url.with_query(None) == auth_config.auth_url assert dict(url.query) == dict( @@ -631,7 +634,10 @@ async def callback(url: URL) -> str: return "test_code" negotiator = HeadlessNegotiator( - connector, config=auth_config, callback=callback, timeout=DEFAULT_TIMEOUT + connector, + config=auth_config, + get_auth_code_cb=get_auth_code_cb, + timeout=DEFAULT_TIMEOUT, ) code = await negotiator.get_code() assert await code.wait() == "test_code" @@ -639,11 +645,14 @@ async def callback(url: URL) -> str: async def test_get_code_raises( self, connector: aiohttp.BaseConnector, auth_config: _AuthConfig ) -> None: - async def callback(url: URL) -> str: + async def get_auth_code_cb(url: URL) -> str: raise RuntimeError("callback error") negotiator = HeadlessNegotiator( - connector, config=auth_config, callback=callback, timeout=DEFAULT_TIMEOUT + connector, + config=auth_config, + get_auth_code_cb=get_auth_code_cb, + timeout=DEFAULT_TIMEOUT, ) with pytest.raises(RuntimeError, match="callback error"): await negotiator.get_code() From 5c96b9481b79a5ec57277e82d78a0f6a99f2d28c Mon Sep 17 00:00:00 2001 From: Andrew Svetlov Date: Tue, 28 May 2019 15:10:12 +0300 Subject: [PATCH 28/28] Fix test --- tests/api/test_config_factory.py | 94 +++++++++++++++++++++----------- 1 file changed, 63 insertions(+), 31 deletions(-) diff --git a/tests/api/test_config_factory.py b/tests/api/test_config_factory.py index 3c776bc90..c53a3efb4 100644 --- a/tests/api/test_config_factory.py +++ b/tests/api/test_config_factory.py @@ -1,6 +1,8 @@ +import abc import sys from pathlib import Path -from typing import Any +from types import TracebackType +from typing import Any, Optional, Type from unittest import mock from uuid import uuid4 as uuid @@ -40,20 +42,38 @@ def config_file( return config_path -@pytest.fixture -async def mock_for_login( - monkeypatch: Any, aiohttp_server: _TestServerFactory -) -> _TestServer: - async def _refresh_token_mock(*args: Any, **kwargs: Any) -> _AuthToken: - return _AuthToken.create_non_expiring(str(uuid())) +class BaseFakeServer(abc.ABC): + def __init__(self): + app = web.Application() + app.router.add_get("/config", self.config_handler) + app.router.add_get("/oauth/show-code", self.show_code) + app.router.add_get("/authorize", self.authorize) + app.router.add_post("/oauth/token", self.token) + self._srv = _TestServer(app) + + @property + def srv(self): + return self._srv + + async def __aenter__(self) -> "BaseFakeServer": + await self._srv.__aenter__() + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: + await self._srv.__aexit__(exc_type, exc_value, traceback) - async def config_handler(request: web.Request) -> web.Response: + async def config_handler(self, request: web.Request) -> web.Response: config_json = { - "auth_url": str(srv.make_url("/authorize")), - "token_url": str(srv.make_url("/oauth/token")), + "auth_url": str(self._srv.make_url("/authorize")), + "token_url": str(self._srv.make_url("/oauth/token")), "client_id": "banana", "audience": "https://test.dev.neuromation.io", - "headless_callback_url": str(srv.make_url("/oauth/show-code")), + "headless_callback_url": str(self._srv.make_url("/oauth/show-code")), "callback_urls": [ "http://127.0.0.2:54540", "http://127.0.0.2:54541", @@ -76,28 +96,42 @@ async def config_handler(request: web.Request) -> web.Response: ) return web.json_response(config_json) - async def show_code(request: web.Request) -> web.Response: - breakpoint() - return web.json_response({}) + @abc.abstractmethod + async def show_code(self, request: web.Request) -> web.Response: + pass - async def authorize(request: web.Request) -> web.Response: - breakpoint() - return web.json_response({}) + @abc.abstractmethod + async def authorize(self, request: web.Request) -> web.Response: + pass + + @abc.abstractmethod + async def token(self, request: web.Request) -> web.Response: + pass - async def token(request: web.Request) -> web.Response: - breakpoint() + +class BrowserFakeServer(BaseFakeServer): + async def show_code(self, request: web.Request) -> web.Response: return web.json_response({}) - app = web.Application() - app.router.add_get("/config", config_handler) - app.router.add_get("/oauth/show-code", show_code) - app.router.add_get("/authorize", authorize) - app.router.add_get("/oauth/token", token) - srv = await aiohttp_server(app) + async def authorize(self, request: web.Request) -> web.Response: + url = URL(request.query['redirect_uri']).with_query(code='test_auth_code') + raise web.HTTPSeeOther(location=url) + + async def token(self, request: web.Request) -> web.Response: + return web.json_response({"access_token": "access_token", + "expires_in": 3600, + "refresh_token": "refresh_token"}) + - monkeypatch.setattr(AuthNegotiator, "refresh_token", _refresh_token_mock) +@pytest.fixture +async def mock_for_login( + monkeypatch: Any, aiohttp_server: _TestServerFactory +) -> _TestServer: + async def _refresh_token_mock(*args: Any, **kwargs: Any) -> _AuthToken: + return _AuthToken.create_non_expiring(str(uuid())) - return srv + async with BrowserFakeServer() as server: + yield server.srv def _create_config( @@ -351,16 +385,14 @@ async def test_normal_login( self, tmp_home: Path, mock_for_login: _TestServer ) -> None: async def get_auth_code_cb(url: URL) -> str: - assert url.with_query(None) == URL( - "https://test-neuromation.auth0.com/authorize" - ) + assert url.with_query(None) == mock_for_login.make_url('/authorize') assert dict(url.query) == dict( response_type="code", code_challenge=mock.ANY, code_challenge_method="S256", client_id="banana", - redirect_uri=mock_for_login.make_url("/oauth/show-code"), + redirect_uri=str(mock_for_login.make_url("/oauth/show-code")), scope="offline_access", audience="https://test.dev.neuromation.io", )