diff --git a/src/main.py b/src/main.py
index 9d105e7..c5a1738 100644
--- a/src/main.py
+++ b/src/main.py
@@ -478,7 +478,6 @@ def display_menu(
     4. Settings
     5. Exit
     """
-    password_manager.update_activity()
     while True:
         if time.time() - password_manager.last_activity > inactivity_timeout:
             print(colored("Session timed out. Vault locked.", "yellow"))
diff --git a/src/tests/test_inactivity_lock.py b/src/tests/test_inactivity_lock.py
index 2958c06..5c6acb1 100644
--- a/src/tests/test_inactivity_lock.py
+++ b/src/tests/test_inactivity_lock.py
@@ -13,6 +13,16 @@ import main
 def test_inactivity_triggers_lock(monkeypatch):
     locked = {"locked": False, "unlocked": False}
 
+    def update_activity():
+        pm.last_activity = time.time()
+
+    def lock_vault():
+        locked["locked"] = True
+
+    def unlock_vault():
+        locked["unlocked"] = True
+        update_activity()
+
     pm = SimpleNamespace(
         is_dirty=False,
         last_update=time.time(),
@@ -21,9 +31,9 @@ def test_inactivity_triggers_lock(monkeypatch):
         handle_add_password=lambda: None,
         handle_retrieve_entry=lambda: None,
         handle_modify_entry=lambda: None,
-        update_activity=lambda: None,
-        lock_vault=lambda: locked.update(locked=True) or None,
-        unlock_vault=lambda: locked.update(unlocked=True) or None,
+        update_activity=update_activity,
+        lock_vault=lock_vault,
+        unlock_vault=unlock_vault,
     )
 
     monkeypatch.setattr("builtins.input", lambda _: "5")