add writePDB ens test

jamesmkrieger · Nov 6, 2023 · 1ac6aa0 · 1ac6aa0
1 parent 39fd57e
commit 1ac6aa0
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/prody/tests/proteins/test_pdbfile.py b/prody/tests/proteins/test_pdbfile.py
@@ -262,6 +262,11 @@ def setUp(self):
         self.ag = parsePDB(self.pdb['path'])
         self.tmp = os.path.join(TEMPDIR, 'test.pdb')
 
+        self.ens = PDBEnsemble()
+        self.ens.setAtoms(self.ag)
+        self.ens.setCoords(self.ag.getCoords())
+        self.ens.addCoordset(self.ag.getCoordsets())
+
         self.ubi = parsePDB(DATA_FILES['1ubi']['path'], secondary=True)
 
         self.hex = parsePDB(DATA_FILES['hex']['path'])
@@ -446,6 +451,16 @@ def testWritingAltlocModels(self):
 
         self.assertEqual(lines1[8], lines2[8],
             'writePDB failed to write correct ANISOU line 8 for 6flr selection with altloc None')
+
+    def testWriteEnsembleToPDB(self):
+        """Test that writePDB can handle ensembles."""
+
+        out = writePDB(self.tmp, self.ens)
+        out = parsePDB(out)
+        self.assertEqual(out.numCoordsets(), self.ens.numCoordsets(),
+            'failed to write correct number of models from ensemble')
+        assert_equal(out.getCoords(), self.ag.getCoordsets(0),
+                'failed to write ensemble model 1 coordinates correctly')
 
     @dec.slow
     def tearDown(self):