From 916813af2d9db3789314103c7d92a14dbb6f6b40 Mon Sep 17 00:00:00 2001
From: Tessa Walsh <tessa@bitarchivist.net>
Date: Mon, 12 Aug 2024 21:51:42 -0400
Subject: [PATCH] Include user and user org info in login response (#2014)

Fixes #2013

Adds the `/users/me` response data to the API login endpoint response
under the key `user_info` and adds a test.
---
 backend/btrixcloud/auth.py | 15 ++++++++++-----
 backend/test/test_users.py | 39 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/backend/btrixcloud/auth.py b/backend/btrixcloud/auth.py
index e08f31a5..81fc27cd 100644
--- a/backend/btrixcloud/auth.py
+++ b/backend/btrixcloud/auth.py
@@ -21,7 +21,7 @@ from fastapi import (
 
 from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
 
-from .models import User
+from .models import User, UserOut
 from .utils import dt_now
 
 
@@ -57,6 +57,7 @@ class BearerResponse(BaseModel):
 
     access_token: str
     token_type: str
+    user_info: UserOut
 
 
 # ============================================================================
@@ -181,10 +182,12 @@ def init_jwt_auth(user_manager):
 
     auth_jwt_router = APIRouter()
 
-    def get_bearer_response(user: User):
+    def get_bearer_response(user: User, user_info: UserOut):
         """get token, return bearer response for user"""
         token = create_access_token(user)
-        return BearerResponse(access_token=token, token_type="bearer")
+        return BearerResponse(
+            access_token=token, token_type="bearer", user_info=user_info
+        )
 
     @auth_jwt_router.post("/login", response_model=BearerResponse)
     async def login(
@@ -246,10 +249,12 @@ def init_jwt_auth(user_manager):
 
         # successfully logged in, reset failed logins, return user
         await user_manager.reset_failed_logins(login_email)
-        return get_bearer_response(user)
+        user_info = await user_manager.get_user_info_with_orgs(user)
+        return get_bearer_response(user, user_info)
 
     @auth_jwt_router.post("/refresh", response_model=BearerResponse)
     async def refresh_jwt(user=Depends(current_active_user)):
-        return get_bearer_response(user)
+        user_info = await user_manager.get_user_info_with_orgs(user)
+        return get_bearer_response(user, user_info)
 
     return auth_jwt_router, current_active_user, shared_secret_or_active_user
diff --git a/backend/test/test_users.py b/backend/test/test_users.py
index c9996469..a7562700 100644
--- a/backend/test/test_users.py
+++ b/backend/test/test_users.py
@@ -78,6 +78,45 @@ def test_me_id(admin_auth_headers, default_org_id):
     assert r.status_code == 404
 
 
+def test_login_user_info(admin_auth_headers, crawler_userid, default_org_id):
+    # Get default org info for comparison
+    r = requests.get(f"{API_PREFIX}/orgs", headers=admin_auth_headers)
+    default_org = [org for org in r.json()["items"] if org["default"]][0]
+
+    # Log in and check response
+    r = requests.post(
+        f"{API_PREFIX}/auth/jwt/login",
+        data={
+            "username": CRAWLER_USERNAME,
+            "password": CRAWLER_PW,
+            "grant_type": "password",
+        },
+    )
+    data = r.json()
+    assert r.status_code == 200
+    assert data["access_token"]
+    assert data["token_type"] == "bearer"
+
+    user_info = data["user_info"]
+    assert user_info
+
+    assert user_info["id"] == crawler_userid
+    assert user_info["name"] == "new-crawler"
+    assert user_info["email"] == CRAWLER_USERNAME
+    assert user_info["is_superuser"] is False
+    assert user_info["is_verified"]
+
+    user_orgs = user_info["orgs"]
+    assert len(user_orgs) == 1
+    org = user_orgs[0]
+
+    assert org["id"] == default_org_id
+    assert org["name"] == default_org["name"]
+    assert org["slug"] == default_org["slug"]
+    assert org["default"]
+    assert org["role"] == 20
+
+
 def test_login_case_insensitive_email():
     r = requests.post(
         f"{API_PREFIX}/auth/jwt/login",