1 files changed, 12 insertions, 15 deletions
diff --git a/gn_auth/auth/authorisation/users/masquerade/views.py b/gn_auth/auth/authorisation/users/masquerade/views.py
index 68f19ee..12a8c97 100644
--- a/gn_auth/auth/authorisation/users/masquerade/views.py
+++ b/gn_auth/auth/authorisation/users/masquerade/views.py
@@ -1,14 +1,14 @@
 """Endpoints for user masquerade"""
 from dataclasses import asdict
 from uuid import UUID
-from functools import partial
 
-from flask import request, jsonify, Response, Blueprint
+from flask import request, jsonify, Response, Blueprint, current_app
 
 from gn_auth.auth.errors import InvalidData
+from gn_auth.auth.authorisation.resources.groups.models import user_group
 
+from ....db import sqlite3 as db
 from ...checks import require_json
-from ....db.sqlite3 import with_db_connection
 from ....authentication.users import user_by_id
 from ....authentication.oauth2.resource_server import require_oauth
 
@@ -21,29 +21,26 @@ masq = Blueprint("masquerade", __name__)
 @require_json
 def masquerade() -> Response:
     """Masquerade as a particular user."""
-    with require_oauth.acquire("profile user masquerade") as token:
+    with (require_oauth.acquire("profile user masquerade") as token,
+          db.connection(current_app.config["AUTH_DB"]) as conn):
         masqueradee_id = UUID(request.json["masquerade_as"])#type: ignore[index]
         if masqueradee_id == token.user.user_id:
             raise InvalidData("You are not allowed to masquerade as yourself.")
 
-        masq_user = with_db_connection(partial(
-            user_by_id, user_id=masqueradee_id))
+        masq_user = user_by_id(conn, user_id=masqueradee_id)
+
         def __masq__(conn):
             new_token = masquerade_as(conn, original_token=token, masqueradee=masq_user)
             return new_token
-        def __dump_token__(tok):
-            return {
-                key: value for key, value in asdict(tok).items()
-                if key in ("access_token", "refresh_token", "expires_in",
-                           "token_type")
-            }
+
         return jsonify({
             "original": {
-                "user": asdict(token.user),
-                "token": __dump_token__(token)
+                "user": asdict(token.user)
             },
             "masquerade_as": {
                 "user": asdict(masq_user),
-                "token": __dump_token__(with_db_connection(__masq__))
+                "token": __masq__(conn),
+                **(user_group(conn, masq_user).maybe(# type: ignore[misc]
+                    {}, lambda grp: {"group": grp}))
             }
         })