Fix tests

closeio · Nov 6, 2024 · 707e821 · 707e821
1 parent 8e26a62
commit 707e821
Showing 1 changed file with 14 additions and 4 deletions.
diff --git a/tests/imap/test_labels.py b/tests/imap/test_labels.py
@@ -48,10 +48,15 @@ def folder_and_message_maps(db, default_account):
 def add_inbox_label(db, default_account, message):
     assert len(message.imapuids) == 1
     imapuid = message.imapuids[0]
-    assert {c.name for c in imapuid.categories} == {"all"}
+    assert {c.name for c in imapuid.get_categories(default_account.category_type)} == {
+        "all"
+    }
     imapuid.update_labels(["\\Inbox"])
     db.session.commit()
-    assert {c.name for c in imapuid.categories} == {"all", "inbox"}
+    assert {c.name for c in imapuid.get_categories(default_account.category_type)} == {
+        "all",
+        "inbox",
+    }
     update_message_metadata(
         db.session, default_account.id, default_account.category_type, message, False
     )
@@ -62,10 +67,15 @@ def add_inbox_label(db, default_account, message):
 def add_custom_label(db, default_account, message):
     assert len(message.imapuids) == 1
     imapuid = message.imapuids[0]
-    existing = [c.name for c in imapuid.categories][0]
+    existing = [c.name for c in imapuid.get_categories(default_account.category_type)][
+        0
+    ]
     imapuid.update_labels(["<3"])
     db.session.commit()
-    assert {c.name for c in imapuid.categories} == {existing, ""}
+    assert {c.name for c in imapuid.get_categories(default_account.category_type)} == {
+        existing,
+        "",
+    }
     update_message_metadata(
         db.session, default_account.id, default_account.category_type, message, False
     )