Refs #253, Added tests for 2FA configuration on login.

teraserver/python/tests/modules/FlaskModule/API/user/test_UserLogin.py

@@ -5,6 +5,41 @@
 class UserLoginTest(BaseUserAPITest):
     test_endpoint = '/api/user/login'
 
+    def setUp(self):
+        super().setUp()
+        # Create users with 2fa enabled
+        with self._flask_app.app_context():
+            self.user1: dict = self._create_2fa_enabled_user('test_user_2fa_1', 'password', set_secret=True)
+            self.user2: dict = self._create_2fa_enabled_user('test_user_2fa_2', 'password', set_secret=False)
+
+    def tearDown(self):
+        # Delete users with 2fa enabled
+        with self._flask_app.app_context():
+            TeraUser.delete(self.user1['id_user'], hard_delete=True)
+            TeraUser.delete(self.user2['id_user'], hard_delete=True)
+        super().tearDown()
+
+    def _create_2fa_enabled_user(self, username, password, set_secret:bool = True):
+        user = TeraUser()
+        user.id_user = 0 # New user
+        user.user_username = username
+        user.user_password = password
+        user.user_firstname = username
+        user.user_lastname = username
+        user.user_email = f"{username}@test.com"
+        user.user_enabled = True
+        user.user_profile = {}
+        if set_secret:
+            user.enable_2fa_otp()
+        else:
+            user.user_2fa_enabled = True
+            user.user_2fa_otp_enabled = False
+            user.user_2fa_otp_secret = None
+
+        TeraUser.insert(user)
+        return user.to_json(minimal=False)
+
+
     def test_get_endpoint_no_auth(self):
         with self._flask_app.app_context():
             response = self.test_client.get(self.test_endpoint)
@@ -62,3 +97,33 @@ def test_get_endpoint_login_admin_user_http_auth_then_token_auth(self):
             # Not allowed for this endpoint
             response = self._get_with_user_token_auth(self.test_client, token=token)
             self.assertEqual(401, response.status_code)
+
+    def test_get_endpoint_login_user1_2fa_already_setup(self):
+        with self._flask_app.app_context():
+
+            # Login should redirect to 2fa verification
+            response = self._get_with_user_http_auth(self.test_client, 'test_user_2fa_1', 'password')
+            self.assertEqual(200, response.status_code)
+            self.assertTrue('redirect_url' in response.json)
+            self.assertFalse('login_setup_2fa' in response.json['redirect_url'])
+            self.assertTrue('login_validate_2fa' in response.json['redirect_url'])
+
+            # Answer should not provide login information
+            self.assertFalse('websocket_url' in response.json)
+            self.assertFalse('user_uuid' in response.json)
+            self.assertFalse('user_token' in response.json)
+
+    def test_get_endpoint_login_user2_2fa_not_setup(self):
+        with self._flask_app.app_context():
+
+            # Login should redirect to 2fa verification
+            response = self._get_with_user_http_auth(self.test_client, 'test_user_2fa_2', 'password')
+            self.assertEqual(200, response.status_code)
+            self.assertTrue('redirect_url' in response.json)
+            self.assertTrue('login_setup_2fa' in response.json['redirect_url'])
+            self.assertFalse('login_validate_2fa' in response.json['redirect_url'])
+
+            # Answer should not provide login information
+            self.assertFalse('websocket_url' in response.json)
+            self.assertFalse('user_uuid' in response.json)
+            self.assertFalse('user_token' in response.json)