[feat] added tests for multiple models

BioPandas · Jul 8, 2024 · cf3a8e8 · cf3a8e8
1 parent 8cc98e5
commit cf3a8e8
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 0 deletions.
diff --git a/biopandas/mmcif/tests/data/2jyf.cif.gz b/biopandas/mmcif/tests/data/2jyf.cif.gz
diff --git a/biopandas/mmcif/tests/test_multiple_models.py b/biopandas/mmcif/tests/test_multiple_models.py
@@ -0,0 +1,26 @@
+# BioPandas
+# Author: Sebastian Raschka <[email protected]>
+# Author: Arian Jamasb <[email protected]>
+# License: BSD 3 clause
+# Project Website: http://rasbt.github.io/biopandas/
+# Code Repository: https://github.com/rasbt/biopandas
+import os
+
+from biopandas.mmcif import PandasMmcif
+
+TESTDATA_FILENAME = os.path.join(os.path.dirname(__file__), "data", "2jyf.cif.gz")
+
+def test_get_model():
+    df = PandasMmcif().read_mmcif(TESTDATA_FILENAME)
+    MODEL_INDEX = 1
+    new_df = df.get_model(MODEL_INDEX)
+    print(df)
+    assert new_df.df["ATOM"]["pdbx_PDB_model_num"].all() == MODEL_INDEX
+
+
+def test_get_models():
+    df = PandasMmcif().read_mmcif(TESTDATA_FILENAME)
+    MODEL_INDICES = [1, 3, 5]
+
+    new_df = df.get_models(MODEL_INDICES)
+    assert new_df.df["ATOM"]["pdbx_PDB_model_num"].all() in MODEL_INDICES