From 1c5ac24e8bcbe8ac7b1329c3c0e39a28b04156df Mon Sep 17 00:00:00 2001
From: Ben White <ben@posthog.com>
Date: Thu, 25 Apr 2024 18:17:40 +0200
Subject: [PATCH] Tried adding tests

---
 posthog/settings/web.py         |  2 +-
 posthog/test/test_middleware.py | 56 +++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/posthog/settings/web.py b/posthog/settings/web.py
index 90697a27ed758..976cdc3ffb4e4 100644
--- a/posthog/settings/web.py
+++ b/posthog/settings/web.py
@@ -353,7 +353,7 @@ def add_recorder_js_headers(headers, path, url):
 PUBLIC_EGRESS_IP_ADDRESSES = get_list(os.getenv("PUBLIC_EGRESS_IP_ADDRESSES", ""))
 
 IMPERSONATION_TIMEOUT_SECONDS = get_from_env("IMPERSONATION_TIMEOUT_SECONDS", 15 * 60, type_cast=int)
-IMPERSONATION_TIMEOUT_SECONDS = 5
+
 # If False, will expire once the session age is greater than IMPERSONATION_TIMEOUT_SECONDS
 # If True, will expire IMPERSONATION_TIMEOUT_SECONDS after the last activity
 IMPERSONATION_EXPIRE_AFTER_LAST_ACTIVITY = get_from_env(
diff --git a/posthog/test/test_middleware.py b/posthog/test/test_middleware.py
index e0f5283dd3cae..3a83abd237f89 100644
--- a/posthog/test/test_middleware.py
+++ b/posthog/test/test_middleware.py
@@ -1,7 +1,12 @@
+from datetime import datetime, timedelta, timezone
 import json
+from time import sleep
 from urllib.parse import quote
 
+from django.conf import settings
 from django.test.client import Client
+from django.urls import reverse
+from freezegun import freeze_time
 from rest_framework import status
 from posthog.api.test.test_organization import create_organization
 from posthog.api.test.test_team import create_team
@@ -9,6 +14,7 @@
 from posthog.models import Action, Cohort, Dashboard, FeatureFlag, Insight
 from posthog.models.organization import Organization
 from posthog.models.team import Team
+from posthog.models.user import User
 from posthog.settings import SITE_URL
 from posthog.test.base import APIBaseTest, override_settings
 
@@ -451,3 +457,53 @@ def test_logout(self):
         # Check if the cookies are not present in the response
         self.assertNotIn("ph_current_project_token", response.cookies)
         self.assertNotIn("ph_current_project_name", response.cookies)
+
+
+@override_settings(IMPERSONATION_TIMEOUT_SECONDS=30)
+class TestAutoLogoutImpersonateMiddleware(APIBaseTest):
+    other_user: User
+
+    def setUp(self):
+        super().setUp()
+        # Reset back to initial team/org for each test
+        self.other_user = User.objects.create_and_join(
+            self.organization, email="other-user@posthog.com", password="123456"
+        )
+
+        self.user.is_staff = True
+        self.user.save()
+
+    def get_csrf_token_payload(self):
+        return {}
+
+    def login_as_other_user(self):
+        return self.client.post(
+            reverse("loginas-user-login", kwargs={"user_id": self.other_user.id}),
+            follow=True,
+        )
+
+    def test_staff_user_can_login(self):
+        assert self.client.get("/api/users/@me").json()["email"] == self.user.email
+        response = self.login_as_other_user()
+        assert response.status_code == 200
+        assert self.client.get("/api/users/@me").json()["email"] == "other-user@posthog.com"
+
+    def test_not_staff_user_cannot_login(self):
+        self.user.is_staff = False
+        self.user.save()
+        assert self.client.get("/api/users/@me").json()["email"] == self.user.email
+        response = self.login_as_other_user()
+        assert response.status_code == 200
+        assert self.client.get("/api/users/@me").json()["email"] == self.user.email
+
+    def test_after_timeout_api_requests_401(self):
+        now = datetime.now()
+        self.login_as_other_user()
+        client = self.client
+        assert client.get("/api/users/@me").status_code == 200
+
+        with freeze_time(now + timedelta(seconds=10)):
+            assert client.get("/api/users/@me").status_code == 200
+
+        with freeze_time(now + timedelta(seconds=35)):
+            assert client.get("/api/users/@me").status_code == 401