Fix and test convenience functions

klieret · Apr 19, 2019 · 3b240b4 · 3b240b4
1 parent c88d6e5
commit 3b240b4
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/ankipandas/convenience_functions.py b/ankipandas/convenience_functions.py
@@ -64,8 +64,8 @@ def load_cards(
     if merge_notes:
         apd.merge_note_info(db, df, inplace=True)
         apd.add_model_names(db, df, inplace=True)
-    if expand_fields:
-        apd.add_fields_as_columns(db, df, inplace=True)
+        if expand_fields:
+            apd.add_fields_as_columns(db, df, inplace=True)
     apd.close_db(db)
     return df
 

diff --git a/ankipandas/test/test_convenience.py b/ankipandas/test/test_convenience.py
@@ -175,5 +175,23 @@ def test_load_notes_expand(self):
             sorted(note_cols + ["mname", "Front", "Back"])
         )
 
+    def test_load_cards(self):
+        cards = convenience.load_cards(self.path)
+        self.assertEqual(
+            sorted(list(cards.columns)),
+            sorted(list(set(
+                card_cols + note_cols + ["dname", "mname", "Front", "Back"] +
+                ["ndata", "nflags", "nmod", "nusn"]
+            )))  # clashes
+        )
+
+    def test_load_cards_nomerge(self):
+        cards = convenience.load_cards(self.path, merge_notes=False)
+        self.assertEqual(
+            sorted(list(cards.columns)),
+            sorted(card_cols + ["dname"])
+        )
+
+
 if __name__ == "__main__":
     unittest.main()