From 791906521f94dcd2545048c1abfb34c47dc6dc9b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@stuvel.eu>
Date: Thu, 30 Aug 2018 18:27:55 +0200
Subject: [PATCH] Added a test context manager to log in when doing Flask test
 client requests

---
 pillar/auth/__init__.py  |  5 +++++
 pillar/tests/__init__.py | 26 ++++++++++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/pillar/auth/__init__.py b/pillar/auth/__init__.py
index aef00a70..68b8db78 100644
--- a/pillar/auth/__init__.py
+++ b/pillar/auth/__init__.py
@@ -210,6 +210,11 @@ def login_user(oauth_token: str, *, load_from_db=False):
         user = _load_user(oauth_token)
     else:
         user = UserClass(oauth_token)
+    login_user_object(user)
+
+
+def login_user_object(user: UserClass):
+    """Log in the given user."""
     flask_login.login_user(user, remember=True)
     g.current_user = user
     user_authenticated.send(None)
diff --git a/pillar/tests/__init__.py b/pillar/tests/__init__.py
index 9e55dac4..b7f2c196 100644
--- a/pillar/tests/__init__.py
+++ b/pillar/tests/__init__.py
@@ -1,6 +1,7 @@
 # -*- encoding: utf-8 -*-
 
 import base64
+import contextlib
 import copy
 import datetime
 import json
@@ -327,6 +328,31 @@ class AbstractPillarTest(TestMinimal):
 
         return user
 
+    @contextlib.contextmanager
+    def login_as(self, user_id: typing.Union[str, ObjectId]):
+        """Context manager, within the context the app context is active and the user logged in.
+
+        The logging-in happens when a request starts, so it's only active when
+        e.g. self.get() or self.post() or somesuch request is used.
+        """
+        from pillar.auth import UserClass, login_user_object
+
+        if isinstance(user_id, str):
+            user_oid = ObjectId(user_id)
+        elif isinstance(user_id, ObjectId):
+            user_oid = user_id
+        else:
+            raise TypeError(f'invalid type {type(user_id)} for parameter user_id')
+        user_doc = self.fetch_user_from_db(user_oid)
+
+        def signal_handler(sender, **kwargs):
+            login_user_object(user)
+
+        with self.app.app_context():
+            user = UserClass.construct('', user_doc)
+            with flask.request_started.connected_to(signal_handler, self.app):
+                yield
+
     def create_valid_auth_token(self, user_id, token='token'):
         from pillar.api.utils import utcnow