modify tests to use CPDB

a-r-j · Feb 6, 2024 · d45de4f · d45de4f
1 parent cc4fff8
commit d45de4f
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/graphein/protein/utils.py b/graphein/protein/utils.py
@@ -493,6 +493,7 @@ def save_rgroup_df_to_pdb(
     :type gz: bool
     """
     ppd = PandasPdb()
+    df = g.graph["rgroup_df"].copy()
 
     # format charge correctly
     df.charge = df.charge.replace("", np.nan)

diff --git a/tests/protein/test_utils.py b/tests/protein/test_utils.py
@@ -28,9 +28,12 @@ def test_save_graph_to_pdb():
 
     # Check for equivalence between saved and existing DFs.
     # We drop the line_idx columns as these will be renumbered
+    #assert_frame_equal(
+    #    a.drop(["line_idx"], axis=1),
+    #    g.graph["pdb_df"].drop(["line_idx", "node_id", "residue_id"], axis=1),
+    #)
     assert_frame_equal(
-        a.drop(["line_idx"], axis=1),
-        g.graph["pdb_df"].drop(["line_idx", "node_id", "residue_id"], axis=1),
+        a, g.graph["pdb_df"].drop(["node_id", "residue_id", axis=1])
     )
     h = construct_graph(path="/tmp/test_graph.pdb")
 
@@ -48,9 +51,12 @@ def test_save_pdb_df_to_pdb():
     assert os.path.isfile("/tmp/test_graph.pdb")
 
     # We drop the line_idx columns as these will be renumbered
+    #assert_frame_equal(
+    #    a.drop(["line_idx"], axis=1),
+    #    g.graph["pdb_df"].drop(["line_idx", "node_id", "residue_id"], axis=1),
+    #)
     assert_frame_equal(
-        a.drop(["line_idx"], axis=1),
-        g.graph["pdb_df"].drop(["line_idx", "node_id", "residue_id"], axis=1),
+        a, g.graph["pdb_df"].drop(["node_id", "residue_id", axis=1])
     )
 
     # Now check for raw, unprocessed DF