Refresh the config once logged in (#766)

CHANGELOG.D/766.feature

@@ -0,0 +1 @@
+Support retrieving server config for authorized users.

neuromation/api/config_factory.py

@@ -54,13 +54,15 @@ async def login(
     ) -> None:
         if self._path.exists():
             raise ConfigError(f"Config file {self._path} already exists. Please logout")
-        server_config = await get_server_config(url)
-        negotiator = AuthNegotiator(server_config.auth_config)
+        config_unauthorized = await get_server_config(url)
+        negotiator = AuthNegotiator(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=server_config.auth_config,
+            auth_config=config_authorized.auth_config,
             auth_token=auth_token,
-            cluster_config=server_config.cluster_config,
+            cluster_config=config_authorized.cluster_config,
             pypi=_PyPIVersion.create_uninitialized(),
             url=url,
         )
@@ -77,7 +79,7 @@ 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)
+        server_config = await get_server_config(url, token=token)
         config = _Config(
             auth_config=server_config.auth_config,
             auth_token=_AuthToken.create_non_expiring(token),

neuromation/api/login.py

@@ -7,7 +7,17 @@
 import time
 import webbrowser
 from dataclasses import dataclass, field
-from typing import Any, AsyncIterator, Callable, List, Optional, Sequence, Type, cast
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Type,
+    cast,
+)
 
 import aiohttp
 from aiohttp import ClientResponseError, ClientSession
@@ -422,38 +432,36 @@ class ConfigLoadException(Exception):
     pass
 
 
-async def get_server_config(url: URL) -> _ServerConfig:
+async def get_server_config(url: URL, token: Optional[str] = None) -> _ServerConfig:
     async with aiohttp.ClientSession(timeout=DEFAULT_TIMEOUT) as client:
-        async with client.get(url / "config") as resp:
+        headers: Dict[str, str] = {}
+        if token:
+            headers["Authorization"] = f"Bearer {token}"
+
+        async with client.get(url / "config", headers=headers) as resp:
             if resp.status != 200:
                 raise RuntimeError(f"Unable to get server configuration: {resp.status}")
             payload = await resp.json()
             # TODO (ajuszkowski, 5-Feb-2019) validate received data
-            auth_url = URL(payload["auth_url"])
-            token_url = URL(payload["token_url"])
-            client_id = payload["client_id"]
-            audience = payload["audience"]
-            success_redirect_url = payload.get("success_redirect_url")
-            if success_redirect_url is not None:
-                success_redirect_url = URL(success_redirect_url)
+            success_redirect_url = URL(payload.get("success_redirect_url", "")) or None
             callback_urls = payload.get("callback_urls")
             callback_urls = (
                 tuple(URL(u) for u in callback_urls)
                 if callback_urls is not None
                 else _AuthConfig.callback_urls
             )
             auth_config = _AuthConfig(
-                auth_url=auth_url,
-                token_url=token_url,
-                client_id=client_id,
-                audience=audience,
+                auth_url=URL(payload["auth_url"]),
+                token_url=URL(payload["token_url"]),
+                client_id=payload["client_id"],
+                audience=payload["audience"],
                 success_redirect_url=success_redirect_url,
                 callback_urls=callback_urls,
             )
-            cluster_config = _ClusterConfig.create(
-                registry_url=URL(payload["registry_url"]),
-                storage_url=URL(payload["storage_url"]),
-                users_url=URL(payload["users_url"]),
-                monitoring_url=URL(payload["monitoring_url"]),
+            cluster_config = _ClusterConfig(
+                registry_url=URL(payload.get("registry_url", "")),
+                storage_url=URL(payload.get("storage_url", "")),
+                users_url=URL(payload.get("users_url", "")),
+                monitoring_url=URL(payload.get("monitoring_url", "")),
             )
-            return _ServerConfig(auth_config=auth_config, cluster_config=cluster_config)
+            return _ServerConfig(cluster_config=cluster_config, auth_config=auth_config)

tests/api/test_config_factory.py

@@ -50,24 +50,29 @@ async def _jobs_list_mock(
         return []
 
     async def _config_handler(request: web.Request) -> web.Response:
-        return web.json_response(
-            {
-                "registry_url": "https://registry-dev.test.com",
-                "storage_url": "https://storage-dev.test.com",
-                "users_url": "https://users-dev.test.com",
-                "monitoring_url": "https://monitoring-dev.test.com",
-                "auth_url": "https://test-neuromation.auth0.com/authorize",
-                "token_url": "https://test-neuromation.auth0.com/oauth/token",
-                "client_id": "banana",
-                "audience": "https://test.dev.neuromation.io",
-                "callback_urls": [
-                    "http://127.0.0.2:54540",
-                    "http://127.0.0.2:54541",
-                    "http://127.0.0.2:54542",
-                ],
-                "success_redirect_url": "https://neu.ro/#test",
-            }
-        )
+        config_json = {
+            "auth_url": "https://test-neuromation.auth0.com/authorize",
+            "token_url": "https://test-neuromation.auth0.com/oauth/token",
+            "client_id": "banana",
+            "audience": "https://test.dev.neuromation.io",
+            "callback_urls": [
+                "http://127.0.0.2:54540",
+                "http://127.0.0.2:54541",
+                "http://127.0.0.2:54542",
+            ],
+            "success_redirect_url": "https://neu.ro/#test",
+        }
+
+        if "Authorization" in request.headers:
+            config_json.update(
+                {
+                    "registry_url": "https://registry-dev.test.com",
+                    "storage_url": "https://storage-dev.test.com",
+                    "users_url": "https://users-dev.test.com",
+                    "monitoring_url": "https://monitoring-dev.test.com",
+                }
+            )
+        return web.json_response(config_json)
 
     app = web.Application()
     app.router.add_get("/config", _config_handler)
@@ -267,7 +272,11 @@ async def test_login_already_logged(self, config_file: Path) -> None:
 
     async def test_normal_login(self, tmp_home: Path, mock_for_login: URL) -> None:
         await Factory().login(url=mock_for_login)
-        assert Path(tmp_home / ".nmrc").exists(), "Config file not written after 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 TestLoginWithToken:
@@ -277,7 +286,11 @@ async def test_login_with_token_already_logged(self, config_file: Path) -> None:
 
     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)
-        assert Path(tmp_home / ".nmrc").exists(), "Config file not written after 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:

tests/api/test_login_utils.py

@@ -13,10 +13,6 @@
 
 
 async def test_get_server_config(aiohttp_server: _TestServerFactory) -> None:
-    registry_url = "https://registry.dev.neuromation.io"
-    storage_url = "https://storage.dev.neuromation.io"
-    users_url = "https://dev.neuromation.io/users"
-    monitoring_url = "https://dev.neuromation.io/monitoring"
     auth_url = "https://dev-neuromation.auth0.com/authorize"
     token_url = "https://dev-neuromation.auth0.com/oauth/token"
     client_id = "this_is_client_id"
@@ -28,10 +24,6 @@ async def test_get_server_config(aiohttp_server: _TestServerFactory) -> None:
     ]
     success_redirect_url = "https://platform.neuromation.io"
     JSON = {
-        "registry_url": registry_url,
-        "storage_url": storage_url,
-        "users_url": users_url,
-        "monitoring_url": monitoring_url,
         "auth_url": auth_url,
         "token_url": token_url,
         "client_id": client_id,
@@ -41,6 +33,7 @@ async def test_get_server_config(aiohttp_server: _TestServerFactory) -> None:
     }
 
     async def handler(request: web.Request) -> web.Response:
+        assert "Authorization" not in request.headers
         return web.json_response(JSON)
 
     app = web.Application()
@@ -58,17 +51,51 @@ async def handler(request: web.Request) -> web.Response:
             success_redirect_url=URL(success_redirect_url),
         ),
         cluster_config=_ClusterConfig.create(
-            registry_url=URL(registry_url),
-            storage_url=URL(storage_url),
-            users_url=URL(users_url),
-            monitoring_url=URL(monitoring_url),
+            registry_url=URL(), storage_url=URL(), users_url=URL(), monitoring_url=URL()
         ),
     )
 
 
 async def test_get_server_config_no_callback_urls(
     aiohttp_server: _TestServerFactory
 ) -> None:
+    auth_url = "https://dev-neuromation.auth0.com/authorize"
+    token_url = "https://dev-neuromation.auth0.com/oauth/token"
+    client_id = "this_is_client_id"
+    audience = "https://platform.dev.neuromation.io"
+    success_redirect_url = "https://platform.neuromation.io"
+    JSON = {
+        "auth_url": auth_url,
+        "token_url": token_url,
+        "client_id": client_id,
+        "audience": audience,
+        "success_redirect_url": success_redirect_url,
+    }
+
+    async def handler(request: web.Request) -> web.Response:
+        assert "Authorization" not in request.headers
+        return web.json_response(JSON)
+
+    app = web.Application()
+    app.router.add_get("/config", handler)
+    srv = await aiohttp_server(app)
+
+    config = await get_server_config(srv.make_url("/"))
+    assert config == _ServerConfig(
+        auth_config=_AuthConfig(
+            auth_url=URL(auth_url),
+            token_url=URL(token_url),
+            client_id=client_id,
+            audience=audience,
+            success_redirect_url=URL(success_redirect_url),
+        ),
+        cluster_config=_ClusterConfig(
+            registry_url=URL(), storage_url=URL(), users_url=URL(), monitoring_url=URL()
+        ),
+    )
+
+
+async def test_get_server_config_with_token(aiohttp_server: _TestServerFactory) -> None:
     registry_url = "https://registry.dev.neuromation.io"
     storage_url = "https://storage.dev.neuromation.io"
     users_url = "https://dev.neuromation.io/users"
@@ -91,13 +118,14 @@ async def test_get_server_config_no_callback_urls(
     }
 
     async def handler(request: web.Request) -> web.Response:
+        assert request.headers["Authorization"] == "Bearer bananatoken"
         return web.json_response(JSON)
 
     app = web.Application()
     app.router.add_get("/config", handler)
     srv = await aiohttp_server(app)
 
-    config = await get_server_config(srv.make_url("/"))
+    config = await get_server_config(srv.make_url("/"), token="bananatoken")
     assert config == _ServerConfig(
         auth_config=_AuthConfig(
             auth_url=URL(auth_url),