diff --git a/app/classes/user.php b/app/classes/user.php
index 94e2756..96427d5 100644
--- a/app/classes/user.php
+++ b/app/classes/user.php
@@ -142,6 +142,8 @@ class User {
             // Login successful
             $_SESSION['user_id'] = $user['id'];
             $_SESSION['username'] = $user['username'];
+            $_SESSION['CREATED'] = time();
+            $_SESSION['LAST_ACTIVITY'] = time();
             return [
                 'status' => 'success',
                 'user_id' => $user['id'],
@@ -151,10 +153,7 @@ class User {
 
         // Get remaining attempts AFTER this failed attempt
         $remainingAttempts = $this->rateLimiter->getRemainingAttempts($username, $ipAddress);
-        return [
-            'status' => 'failed',
-            'message' => "Invalid credentials. {$remainingAttempts} attempts remaining."
-        ];
+        throw new Exception("Invalid credentials. {$remainingAttempts} attempts remaining.");
     }
 
 
diff --git a/tests/Feature/Middleware/SessionMiddlewareTest.php b/tests/Feature/Middleware/SessionMiddlewareTest.php
index 171c9d4..cf472eb 100644
--- a/tests/Feature/Middleware/SessionMiddlewareTest.php
+++ b/tests/Feature/Middleware/SessionMiddlewareTest.php
@@ -24,6 +24,14 @@ class SessionMiddlewareTest extends TestCase
             'domain' => 'localhost'
         ];
         $this->app_root = 'https://localhost/app';
+
+        // Initialize session variables
+        $_SESSION = [
+            'user_id' => 1,
+            'username' => 'testuser',
+            'CREATED' => time(),
+            'LAST_ACTIVITY' => time()
+        ];
     }
 
     protected function tearDown(): void
@@ -33,78 +41,65 @@ class SessionMiddlewareTest extends TestCase
 
     public function testSessionStart()
     {
-        $_SESSION = ['USER_ID' => 1];
         $result = applySessionMiddleware($this->config, $this->app_root);
 
         $this->assertTrue($result);
         $this->assertArrayHasKey('LAST_ACTIVITY', $_SESSION);
         $this->assertArrayHasKey('CREATED', $_SESSION);
-        $this->assertArrayHasKey('USER_ID', $_SESSION);
-        $this->assertEquals(1, $_SESSION['USER_ID']);
+        $this->assertArrayHasKey('user_id', $_SESSION);
+        $this->assertEquals(1, $_SESSION['user_id']);
     }
 
     public function testSessionTimeout()
     {
-        $_SESSION = [
-            'USER_ID' => 1,
-            'LAST_ACTIVITY' => time() - 1500 // 25 minutes ago
-        ];
+        $_SESSION['LAST_ACTIVITY'] = time() - 1500; // 25 minutes ago
 
         $result = applySessionMiddleware($this->config, $this->app_root);
 
         $this->assertFalse($result);
-        $this->assertArrayNotHasKey('USER_ID', $_SESSION, 'Session should be cleared after timeout');
+        $this->assertArrayNotHasKey('user_id', $_SESSION, 'Session should be cleared after timeout');
     }
 
     public function testSessionRegeneration()
     {
         $now = time();
-        $_SESSION = [
-            'USER_ID' => 1,
-            'CREATED' => $now - 1900 // 31+ minutes ago
-        ];
+        $_SESSION['CREATED'] = $now - 1900; // 31+ minutes ago
 
         $result = applySessionMiddleware($this->config, $this->app_root);
 
         $this->assertTrue($result);
-        $this->assertEquals(1, $_SESSION['USER_ID']);
+        $this->assertEquals(1, $_SESSION['user_id']);
         $this->assertGreaterThanOrEqual($now - 1900, $_SESSION['CREATED']);
         $this->assertLessThanOrEqual($now + 10, $_SESSION['CREATED']);
     }
 
     public function testRememberMe()
     {
-        $_SESSION = [
-            'USER_ID' => 1,
-            'REMEMBER_ME' => true,
-            'LAST_ACTIVITY' => time() - 86500 // More than 24 hours ago
-        ];
+        $_SESSION['REMEMBER_ME'] = true;
+        $_SESSION['LAST_ACTIVITY'] = time() - 86500; // More than 24 hours ago
 
         $result = applySessionMiddleware($this->config, $this->app_root);
 
         $this->assertTrue($result);
-        $this->assertArrayHasKey('USER_ID', $_SESSION);
+        $this->assertArrayHasKey('user_id', $_SESSION);
     }
 
     public function testNoUserSession()
     {
-        $_SESSION = [];
+        unset($_SESSION['user_id']);
         $result = applySessionMiddleware($this->config, $this->app_root);
 
         $this->assertFalse($result);
-        $this->assertArrayNotHasKey('USER_ID', $_SESSION);
+        $this->assertArrayNotHasKey('user_id', $_SESSION);
     }
 
     public function testSessionHeaders()
     {
-        $_SESSION = [
-            'USER_ID' => 1,
-            'LAST_ACTIVITY' => time() - 1500 // 25 minutes ago
-        ];
+        $_SESSION['LAST_ACTIVITY'] = time() - 1500; // 25 minutes ago
 
         $result = applySessionMiddleware($this->config, $this->app_root);
 
         $this->assertFalse($result);
-        $this->assertArrayNotHasKey('USER_ID', $_SESSION, 'Session should be cleared after timeout');
+        $this->assertArrayNotHasKey('user_id', $_SESSION, 'Session should be cleared after timeout');
     }
 }
diff --git a/tests/Unit/Classes/UserTest.php b/tests/Unit/Classes/UserTest.php
index 672f887..13ff261 100644
--- a/tests/Unit/Classes/UserTest.php
+++ b/tests/Unit/Classes/UserTest.php
@@ -44,6 +44,19 @@ class UserTest extends TestCase
             )
         ");
 
+        // Create user_2fa table for two-factor authentication
+        $this->db->getConnection()->exec("
+            CREATE TABLE user_2fa (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                user_id INTEGER NOT NULL,
+                secret_key TEXT NOT NULL,
+                backup_codes TEXT,
+                enabled TINYINT(1) NOT NULL DEFAULT 0,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
+            )
+        ");
+
         // Create tables for rate limiter
         $this->db->getConnection()->exec("
             CREATE TABLE login_attempts (
@@ -116,7 +129,13 @@ class UserTest extends TestCase
         // Test successful login
         try {
             $result = $this->user->login('testuser', $password);
-            $this->assertTrue($result);
+            $this->assertIsArray($result);
+            $this->assertEquals('success', $result['status']);
+            $this->assertArrayHasKey('user_id', $result);
+            $this->assertArrayHasKey('username', $result);
+            $this->assertArrayHasKey('user_id', $_SESSION);
+            $this->assertArrayHasKey('CREATED', $_SESSION);
+            $this->assertArrayHasKey('LAST_ACTIVITY', $_SESSION);
         } catch (Exception $e) {
             $this->fail('Login should not throw an exception for valid credentials: ' . $e->getMessage());
         }