From 65b37a4010d8b4cc630d0dad3598db9855ebbf25 Mon Sep 17 00:00:00 2001
From: Cunliang Geng <c.geng@esciencecenter.nl>
Date: Tue, 2 Jun 2020 14:10:06 +0200
Subject: [PATCH] add test for aug_data method

---
 test/test_generate.py | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/test/test_generate.py b/test/test_generate.py
index 5a83a91e..3d7990e7 100644
--- a/test/test_generate.py
+++ b/test/test_generate.py
@@ -257,6 +257,28 @@ def test_7_realign(self):
         database = DataGenerator(hdf5=copy_name)
         database.realign_complexes(align={'axis':'z'})
 
+    def test_8_aug_data(self):
+
+        src_name = './1ak4.hdf5'
+        copy_name = './1ak4_aug.hdf5'
+
+        shutil.copy(src_name, copy_name)
+
+        database = DataGenerator(hdf5=copy_name)
+        database.aug_data(augmentation=2, keep_existing_aug=True)
+        grid_info = {
+            'number_of_points': [10, 10, 10],
+            'resolution': [3., 3., 3.],
+            'atomic_densities': {'C': 1.7, 'N': 1.55, 'O': 1.52, 'S': 1.8},
+            'feature': ['PSSM_ALA','RCD_total', 'bsa', 'charge',],
+        }
+        database.map_features(
+            grid_info,
+            try_sparse=True,
+            time=False,
+            prog_bar=False,
+        )
+
 if __name__ == "__main__":
 
     # unittest.main()
@@ -267,4 +289,5 @@ def test_7_realign(self):
     inst.test_4_add_feature()
     inst.test_5_align()
     inst.test_6_align_interface()
-    inst.test_7_realign()
\ No newline at end of file
+    inst.test_7_realign()
+    inst.test_8_aug_data()
\ No newline at end of file