From a755314dc8320366dd6c369da2982789c6989349 Mon Sep 17 00:00:00 2001
From: Andrew Baldwin <andrew.baldwin44@gmail.com>
Date: Tue, 9 Jan 2024 19:34:55 -0500
Subject: [PATCH] Add web_auth tests

---
 locust/test/test_web.py | 51 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/locust/test/test_web.py b/locust/test/test_web.py
index 133e927b9d..0670778b92 100644
--- a/locust/test/test_web.py
+++ b/locust/test/test_web.py
@@ -14,6 +14,7 @@
 import requests
 from pyquery import PyQuery as pq
 import locust
+from flask_login import UserMixin
 from locust import constant, LoadTestShape
 from locust.argument_parser import get_parser, parse_options
 from locust.user import User, task
@@ -1064,6 +1065,56 @@ def test_logs(self):
         self.assertIn(log_line, response.json().get("logs"))
 
 
+class TestWebUIAuth(LocustTestCase):
+    def setUp(self):
+        super().setUp()
+
+        parser = get_parser(default_config_files=[])
+        self.environment.parsed_options = parser.parse_args(["--modern-ui", "--web-auth"])
+
+        self.web_ui = self.environment.create_web_ui("127.0.0.1", 0, modern_ui=True, use_auth=True)
+
+        self.web_ui.app.secret_key = "secret!"
+        gevent.sleep(0.01)
+        self.web_port = self.web_ui.server.server_port
+
+    def tearDown(self):
+        super().tearDown()
+        self.web_ui.stop()
+        self.runner.quit()
+
+    def test_index_with_web_auth_enabled_valid_user(self):
+        class User(UserMixin):
+            def __init__(self):
+                self.username = "test_user"
+
+            def get_id(self):
+                return self.username
+
+        def load_user(id):
+            return User()
+
+        self.web_ui.login_manager.request_loader(load_user)
+
+        response = requests.get("http://127.0.0.1:%i" % self.web_port)
+        d = pq(response.content.decode("utf-8"))
+
+        self.assertNotIn("authArgs", str(d))
+        self.assertIn("templateArgs", str(d))
+
+    def test_index_with_web_auth_enabled_no_user(self):
+        def load_user():
+            return None
+
+        self.web_ui.login_manager.user_loader(load_user)
+
+        response = requests.get("http://127.0.0.1:%i" % self.web_port)
+        d = pq(response.content.decode("utf-8"))
+
+        # asserts auth page is returned
+        self.assertIn("authArgs", str(d))
+
+
 class TestWebUIWithTLS(LocustTestCase):
     def setUp(self):
         super().setUp()