From f2b6219dc0008f0193d83cecf77b23f3c623eb4d Mon Sep 17 00:00:00 2001
From: thePR0M3TH3AN <53631862+PR0M3TH3AN@users.noreply.github.com>
Date: Fri, 4 Jul 2025 19:31:31 -0400
Subject: [PATCH] Add schema v3 migration and update tests

---
 src/password_manager/migrations.py    | 13 +++++++++-
 src/tests/test_backup_restore.py      | 14 ++++++----
 src/tests/test_index_import_export.py | 22 +++++++++++++---
 src/tests/test_migrations.py          | 37 ++++++++++++++++++++++++---
 4 files changed, 73 insertions(+), 13 deletions(-)

diff --git a/src/password_manager/migrations.py b/src/password_manager/migrations.py
index 3e20295..dbcb16a 100644
--- a/src/password_manager/migrations.py
+++ b/src/password_manager/migrations.py
@@ -39,7 +39,18 @@ def _v1_to_v2(data: dict) -> dict:
     return data
 
 
-LATEST_VERSION = 2
+@migration(2)
+def _v2_to_v3(data: dict) -> dict:
+    """Add custom_fields and origin defaults to each entry."""
+    entries = data.get("entries", {})
+    for entry in entries.values():
+        entry.setdefault("custom_fields", [])
+        entry.setdefault("origin", "")
+    data["schema_version"] = 3
+    return data
+
+
+LATEST_VERSION = 3
 
 
 def apply_migrations(data: dict) -> dict:
diff --git a/src/tests/test_backup_restore.py b/src/tests/test_backup_restore.py
index 312f14a..1c349e1 100644
--- a/src/tests/test_backup_restore.py
+++ b/src/tests/test_backup_restore.py
@@ -22,7 +22,7 @@ def test_backup_restore_workflow(monkeypatch):
         index_file = fp_dir / "seedpass_entries_db.json.enc"
 
         data1 = {
-            "schema_version": 2,
+            "schema_version": 3,
             "entries": {
                 "0": {
                     "website": "a",
@@ -30,6 +30,8 @@ def test_backup_restore_workflow(monkeypatch):
                     "type": "password",
                     "kind": "password",
                     "notes": "",
+                    "custom_fields": [],
+                    "origin": "",
                 }
             },
         }
@@ -43,7 +45,7 @@ def test_backup_restore_workflow(monkeypatch):
         assert backup1.stat().st_mode & 0o777 == 0o600
 
         data2 = {
-            "schema_version": 2,
+            "schema_version": 3,
             "entries": {
                 "0": {
                     "website": "b",
@@ -51,6 +53,8 @@ def test_backup_restore_workflow(monkeypatch):
                     "type": "password",
                     "kind": "password",
                     "notes": "",
+                    "custom_fields": [],
+                    "origin": "",
                 }
             },
         }
@@ -63,11 +67,11 @@ def test_backup_restore_workflow(monkeypatch):
         assert backup2.exists()
         assert backup2.stat().st_mode & 0o777 == 0o600
 
-        vault.save_index({"schema_version": 2, "entries": {"temp": {}}})
+        vault.save_index({"schema_version": 3, "entries": {"temp": {}}})
         backup_mgr.restore_latest_backup()
         assert vault.load_index()["entries"] == data2["entries"]
 
-        vault.save_index({"schema_version": 2, "entries": {}})
+        vault.save_index({"schema_version": 3, "entries": {}})
         backup_mgr.restore_backup_by_timestamp(1111)
         assert vault.load_index()["entries"] == data1["entries"]
 
@@ -85,7 +89,7 @@ def test_additional_backup_location(monkeypatch):
         cfg_mgr.set_additional_backup_path(extra)
         backup_mgr = BackupManager(fp_dir, cfg_mgr)
 
-        vault.save_index({"schema_version": 2, "entries": {"a": {}}})
+        vault.save_index({"schema_version": 3, "entries": {"a": {}}})
 
         monkeypatch.setattr(time, "time", lambda: 3333)
         backup_mgr.create_backup()
diff --git a/src/tests/test_index_import_export.py b/src/tests/test_index_import_export.py
index fce78c2..5742cf3 100644
--- a/src/tests/test_index_import_export.py
+++ b/src/tests/test_index_import_export.py
@@ -31,8 +31,16 @@ def test_index_export_import_round_trip():
         vault = setup_vault(tmp)
 
         original = {
-            "schema_version": 2,
-            "entries": {"0": {"website": "example", "type": "password", "notes": ""}},
+            "schema_version": 3,
+            "entries": {
+                "0": {
+                    "website": "example",
+                    "type": "password",
+                    "notes": "",
+                    "custom_fields": [],
+                    "origin": "",
+                }
+            },
         }
         vault.save_index(original)
 
@@ -41,9 +49,15 @@ def test_index_export_import_round_trip():
 
         vault.save_index(
             {
-                "schema_version": 2,
+                "schema_version": 3,
                 "entries": {
-                    "0": {"website": "changed", "type": "password", "notes": ""}
+                    "0": {
+                        "website": "changed",
+                        "type": "password",
+                        "notes": "",
+                        "custom_fields": [],
+                        "origin": "",
+                    }
                 },
             }
         )
diff --git a/src/tests/test_migrations.py b/src/tests/test_migrations.py
index a1cd7e4..87106cc 100644
--- a/src/tests/test_migrations.py
+++ b/src/tests/test_migrations.py
@@ -13,17 +13,24 @@ def setup(tmp_path: Path):
     return enc_mgr, vault
 
 
-def test_migrate_v0_to_v2(tmp_path: Path):
+def test_migrate_v0_to_v3(tmp_path: Path):
     enc_mgr, vault = setup(tmp_path)
     legacy = {"passwords": {"0": {"website": "a", "length": 8}}}
     enc_mgr.save_json_data(legacy)
     data = vault.load_index()
     assert data["schema_version"] == LATEST_VERSION
-    expected_entry = {"website": "a", "length": 8, "type": "password", "notes": ""}
+    expected_entry = {
+        "website": "a",
+        "length": 8,
+        "type": "password",
+        "notes": "",
+        "custom_fields": [],
+        "origin": "",
+    }
     assert data["entries"]["0"] == expected_entry
 
 
-def test_migrate_v1_to_v2(tmp_path: Path):
+def test_migrate_v1_to_v3(tmp_path: Path):
     enc_mgr, vault = setup(tmp_path)
     legacy = {"schema_version": 1, "passwords": {"0": {"website": "b", "length": 10}}}
     enc_mgr.save_json_data(legacy)
@@ -34,6 +41,30 @@ def test_migrate_v1_to_v2(tmp_path: Path):
         "length": 10,
         "type": "password",
         "notes": "",
+        "custom_fields": [],
+        "origin": "",
+    }
+    assert data["entries"]["0"] == expected_entry
+
+
+def test_migrate_v2_to_v3(tmp_path: Path):
+    enc_mgr, vault = setup(tmp_path)
+    legacy = {
+        "schema_version": 2,
+        "entries": {
+            "0": {"website": "c", "length": 5, "type": "password", "notes": ""}
+        },
+    }
+    enc_mgr.save_json_data(legacy)
+    data = vault.load_index()
+    assert data["schema_version"] == LATEST_VERSION
+    expected_entry = {
+        "website": "c",
+        "length": 5,
+        "type": "password",
+        "notes": "",
+        "custom_fields": [],
+        "origin": "",
     }
     assert data["entries"]["0"] == expected_entry