fix column drops

a-r-j · Feb 7, 2024 · 9f4eda1 · 9f4eda1
1 parent a73e13c
commit 9f4eda1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/protein/test_utils.py b/tests/protein/test_utils.py
@@ -33,7 +33,7 @@ def test_save_graph_to_pdb():
     #    g.graph["pdb_df"].drop(["line_idx", "node_id", "residue_id"], axis=1),
     # )
     assert_frame_equal(
-        a, g.graph["pdb_df"].drop(["node_id", "residue_id"], axis=1)
+        a, g.graph["pdb_df"].drop(["node_id", "residue_id", "line_idx", "blank_1", "blank_2", "blank_3", "blank_4", "segment_id"], axis=1)
     )
     h = construct_graph(path="/tmp/test_graph.pdb")
 
@@ -56,7 +56,7 @@ def test_save_pdb_df_to_pdb():
     #    g.graph["pdb_df"].drop(["line_idx", "node_id", "residue_id"], axis=1),
     # )
     assert_frame_equal(
-        a, g.graph["pdb_df"].drop(["node_id", "residue_id"], axis=1)
+        a, g.graph["pdb_df"].drop(["node_id", "residue_id", "line_idx", "blank_1", "blank_2", "blank_3", "blank_4", "segment_id"], axis=1)
     )
 
     # Now check for raw, unprocessed DF