Merge pull request #145 from lsst-sqre/tickets/DM-34613

[DM-34613] Make UID optional for users
lsst-sqre · Jun 14, 2022 · a5f907c · a5f907c
2 parents e2f401c + 78368b2
commit a5f907c
Show file tree

Hide file tree

Showing 12 changed files with 67 additions and 58 deletions.
diff --git a/src/mobu/flock.py b/src/mobu/flock.py
@@ -133,6 +133,10 @@ def _users_from_spec(self, spec: UserSpec, count: int) -> List[User]:
         users = []
         for i in range(1, count + 1):
             username = spec.username_prefix + str(i).zfill(padding)
-            user = User(username=username, uidnumber=spec.uid_start + i - 1)
+            if spec.uid_start is not None:
+                uid = spec.uid_start + i - 1
+            else:
+                uid = None
+            user = User(username=username, uidnumber=uid)
             users.append(user)
         return users
diff --git a/src/mobu/main.py b/src/mobu/main.py
@@ -47,12 +47,14 @@
 app.include_router(internal_router)
 app.include_router(external_router, prefix=f"/{config.name}")
 
+# Add middleware.
+app.add_middleware(XForwardedMiddleware)
+
 
 @app.on_event("startup")
 async def startup_event() -> None:
     if not config.environment_url:
         raise RuntimeError("ENVIRONMENT_URL was not set")
-    app.add_middleware(XForwardedMiddleware)
     await monkey_business_manager.init()
     if config.autostart:
         await autostart()

diff --git a/src/mobu/models/user.py b/src/mobu/models/user.py
@@ -1,7 +1,7 @@
 """Data models for an authenticated user."""
 
 import time
-from typing import List
+from typing import List, Optional
 
 from aiohttp import ClientSession
 from pydantic import BaseModel, Field
@@ -16,7 +16,15 @@ class User(BaseModel):
 
     username: str = Field(..., title="Username", example="testuser")
 
-    uidnumber: int = Field(..., title="Numeric UID", example=60001)
+    uidnumber: Optional[int] = Field(
+        None,
+        title="Numeric UID",
+        description=(
+            "If omitted, Gafaelfawr will assign a UID. (Gafaelfawr UID"
+            " assignment requires Firestore be configured.)"
+        ),
+        example=60001,
+    )
 
 
 class UserSpec(BaseModel):
@@ -29,10 +37,14 @@ class UserSpec(BaseModel):
         example="lsptestuser",
     )
 
-    uid_start: int = Field(
-        ...,
+    uid_start: Optional[int] = Field(
+        None,
         title="Starting UID",
-        description="Users will be given consecutive UIDs starting with this",
+        description=(
+            "Users will be given consecutive UIDs starting with this. If"
+            " omitted, Gafaelfawr will assign UIDs. (Gafaelfawr UID assignment"
+            " requires Firestore be configured.)"
+        ),
         example=60000,
     )
 
@@ -57,18 +69,20 @@ async def create(
         cls, user: User, scopes: List[str], session: ClientSession
     ) -> "AuthenticatedUser":
         token_url = f"{config.environment_url}/auth/api/v1/tokens"
+        data = {
+            "username": user.username,
+            "name": "Mobu Test User",
+            "token_type": "user",
+            "token_name": f"mobu {str(float(time.time()))}",
+            "scopes": scopes,
+            "expires": int(time.time() + 60 * 60 * 24 * 365),
+        }
+        if user.uidnumber is not None:
+            data["uid"] = user.uidnumber
         r = await session.post(
             token_url,
             headers={"Authorization": f"Bearer {config.gafaelfawr_token}"},
-            json={
-                "username": user.username,
-                "name": "Mobu Test User",
-                "token_type": "user",
-                "token_name": f"mobu {str(float(time.time()))}",
-                "scopes": scopes,
-                "expires": int(time.time() + 60 * 60 * 24 * 365),
-                "uid": user.uidnumber,
-            },
+            json=data,
             raise_for_status=True,
         )
         body = await r.json()

diff --git a/tests/autostart_test.py b/tests/autostart_test.py
@@ -44,7 +44,7 @@ def configure_autostart(
     tmp_path: Path, mock_aioresponses: aioresponses
 ) -> Iterator[None]:
     """Set up the autostart configuration."""
-    mock_gafaelfawr(mock_aioresponses)
+    mock_gafaelfawr(mock_aioresponses, any_uid=True)
     autostart_path = tmp_path / "autostart.yaml"
     autostart_path.write_text(AUTOSTART_CONFIG)
     config.autostart = str(autostart_path)

diff --git a/tests/business/jupyterloginloop_test.py b/tests/business/jupyterloginloop_test.py
@@ -28,7 +28,7 @@ async def test_run(
         json={
             "name": "test",
             "count": 1,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "options": {
                 "spawn_settle_time": 0,
@@ -55,7 +55,6 @@ async def test_run(
         "user": {
             "scopes": ["exec:notebook"],
             "token": ANY,
-            "uidnumber": 1000,
             "username": "testuser1",
         },
     }
@@ -87,7 +86,7 @@ async def test_reuse_lab(
         json={
             "name": "test",
             "count": 1,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "options": {
                 "spawn_settle_time": 0,
@@ -119,7 +118,7 @@ async def test_delayed_lab_delete(
         json={
             "name": "test",
             "count": 5,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "options": {
                 "spawn_settle_time": 0,
@@ -154,7 +153,7 @@ async def test_alert(
         json={
             "name": "test",
             "count": 2,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "options": {
                 "spawn_settle_time": 0,
@@ -215,7 +214,7 @@ async def test_redirect_loop(
         json={
             "name": "test",
             "count": 1,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "options": {
                 "spawn_settle_time": 0,
@@ -281,7 +280,7 @@ async def test_spawn_timeout(
         json={
             "name": "test",
             "count": 1,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "options": {
                 "spawn_settle_time": 0,
@@ -338,7 +337,7 @@ async def test_spawn_failed(
         json={
             "name": "test",
             "count": 1,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "options": {
                 "spawn_settle_time": 0,
@@ -409,7 +408,7 @@ async def test_delete_timeout(
         json={
             "name": "test",
             "count": 1,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "options": {
                 "spawn_settle_time": 0,

diff --git a/tests/business/jupyterpythonloop_test.py b/tests/business/jupyterpythonloop_test.py
@@ -18,7 +18,7 @@
 async def test_run(
     client: AsyncClient, mock_aioresponses: aioresponses
 ) -> None:
-    mock_gafaelfawr(mock_aioresponses)
+    mock_gafaelfawr(mock_aioresponses, username="testuser1", uid=1000)
 
     r = await client.put(
         "/mobu/flocks",
@@ -77,7 +77,7 @@ async def test_server_shutdown(
         json={
             "name": "test",
             "count": 20,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "options": {
                 "spawn_settle_time": 0,
@@ -107,7 +107,7 @@ async def test_alert(
         json={
             "name": "test",
             "count": 1,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "options": {
                 "code": 'raise Exception("some error")',
@@ -193,7 +193,7 @@ async def test_long_error(
         json={
             "name": "test",
             "count": 1,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "options": {
                 "code": "long_error_for_test()",

diff --git a/tests/business/notebookrunner_test.py b/tests/business/notebookrunner_test.py
@@ -41,7 +41,7 @@ async def test_run(
         json={
             "name": "test",
             "count": 1,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "options": {
                 "spawn_settle_time": 0,
@@ -73,7 +73,6 @@ async def test_run(
         "user": {
             "scopes": ["exec:notebook"],
             "token": ANY,
-            "uidnumber": 1000,
             "username": "testuser1",
         },
     }
@@ -117,7 +116,7 @@ async def test_alert(
         json={
             "name": "test",
             "count": 1,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "options": {
                 "spawn_settle_time": 0,
@@ -157,7 +156,6 @@ async def test_alert(
         "user": {
             "scopes": ["exec:notebook"],
             "token": ANY,
-            "uidnumber": 1000,
             "username": "testuser1",
         },
     }

diff --git a/tests/business/tapqueryrunner_test.py b/tests/business/tapqueryrunner_test.py
@@ -34,10 +34,7 @@ async def test_run(
             json={
                 "name": "test",
                 "count": 1,
-                "user_spec": {
-                    "username_prefix": "testuser",
-                    "uid_start": 1000,
-                },
+                "user_spec": {"username_prefix": "testuser"},
                 "scopes": ["exec:notebook"],
                 "business": "TAPQueryRunner",
             },
@@ -59,7 +56,6 @@ async def test_run(
             "user": {
                 "scopes": ["exec:notebook"],
                 "token": ANY,
-                "uidnumber": 1000,
                 "username": "testuser1",
             },
         }
@@ -85,10 +81,7 @@ async def test_alert(
             json={
                 "name": "test",
                 "count": 1,
-                "user_spec": {
-                    "username_prefix": "testuser",
-                    "uid_start": 1000,
-                },
+                "user_spec": {"username_prefix": "testuser"},
                 "scopes": ["exec:notebook"],
                 "business": "TAPQueryRunner",
             },
@@ -161,7 +154,7 @@ async def test_random_object() -> None:
 
     logger = structlog.get_logger(__file__)
     user = AuthenticatedUser(
-        username="user", uidnumber=1000, scopes=["read:tap"], token="blah blah"
+        username="user", scopes=["read:tap"], token="blah blah"
     )
     with patch.object(pyvo.dal, "TAPService"):
         runner = TAPQueryRunner(logger, BusinessConfig(), user)

diff --git a/tests/handlers/flock_test.py b/tests/handlers/flock_test.py
@@ -22,7 +22,7 @@ async def test_start_stop(
     config = {
         "name": "test",
         "count": 1,
-        "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+        "user_spec": {"username_prefix": "testuser"},
         "scopes": ["exec:notebook"],
         "business": "Business",
     }
@@ -33,7 +33,7 @@ async def test_start_stop(
         "config": {
             "name": "test",
             "count": 1,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "business": "Business",
         },
@@ -51,7 +51,6 @@ async def test_start_stop(
                 "user": {
                     "scopes": ["exec:notebook"],
                     "token": ANY,
-                    "uidnumber": 1000,
                     "username": "testuser1",
                 },
             },
@@ -132,7 +131,7 @@ async def test_start_stop(
 async def test_user_list(
     client: AsyncClient, mock_aioresponses: aioresponses
 ) -> None:
-    mock_gafaelfawr(mock_aioresponses)
+    mock_gafaelfawr(mock_aioresponses, any_uid=True)
 
     config = {
         "name": "test",

diff --git a/tests/monkeyflocker_test.py b/tests/monkeyflocker_test.py
@@ -66,7 +66,6 @@ async def startup_event() -> None:
 count: 1
 user_spec:
   username_prefix: testuser
-  uid_start: 1000
 scopes: ["exec:notebook"]
 business: Business
 """
@@ -154,7 +153,7 @@ def test_start_report_stop(tmp_path: Path, app_url: str) -> None:
         "config": {
             "name": "basic",
             "count": 1,
-            "user_spec": {"username_prefix": "testuser", "uid_start": 1000},
+            "user_spec": {"username_prefix": "testuser"},
             "scopes": ["exec:notebook"],
             "business": "Business",
         },
@@ -172,7 +171,6 @@ def test_start_report_stop(tmp_path: Path, app_url: str) -> None:
                 "user": {
                     "scopes": ["exec:notebook"],
                     "token": ANY,
-                    "uidnumber": 1000,
                     "username": "testuser1",
                 },
             },

diff --git a/tests/support/gafaelfawr.py b/tests/support/gafaelfawr.py
@@ -35,6 +35,8 @@ def mock_gafaelfawr(
     mocked: aioresponses,
     username: Optional[str] = None,
     uid: Optional[int] = None,
+    *,
+    any_uid: bool = False,
 ) -> None:
     """Mock out the call to Gafaelfawr to create a user token.
 
@@ -47,19 +49,19 @@ def mock_gafaelfawr(
 
     def handler(url: str, **kwargs: Any) -> CallbackResult:
         assert kwargs["headers"] == {"Authorization": f"Bearer {admin_token}"}
-        assert kwargs["json"] == {
-            "username": ANY,
+        expected = {
+            "username": username if username else ANY,
             "token_type": "user",
             "token_name": ANY,
             "scopes": ["exec:notebook"],
             "expires": ANY,
             "name": "Mobu Test User",
-            "uid": ANY,
         }
-        if username:
-            assert kwargs["json"]["username"] == username
         if uid:
-            assert kwargs["json"]["uid"] == uid
+            expected["uid"] = uid
+        elif any_uid:
+            expected["uid"] = ANY
+        assert kwargs["json"] == expected
         assert kwargs["json"]["token_name"].startswith("mobu ")
         assert kwargs["json"]["expires"] > time.time()
         response = {"token": make_gafaelfawr_token(kwargs["json"]["username"])}