Update tests.

arnor-sigurdsson · Aug 27, 2023 · 68f31fb · 68f31fb
1 parent e67f3db
commit 68f31fb
Showing 1 changed file with 15 additions and 1 deletion.
diff --git a/tests/test_data_load/test_data_preparation_modules/test_prepare_omics.py b/tests/test_data_load/test_data_preparation_modules/test_prepare_omics.py
@@ -19,7 +19,7 @@ def test_prepare_genotype_array_train_mode():
         na_augment_perc=1.0,
         na_augment_prob=1.0,
         shuffle_augment_perc=1.0,
-        shuffle_augment_prob=1.0,
+        shuffle_augment_prob=0.0,
         test_mode=False,
     )
 
@@ -28,6 +28,20 @@ def test_prepare_genotype_array_train_mode():
 
     assert (prepared_array_train[:, -1, :] == 1).all()
 
+    prepared_array_train = prepare_omics.prepare_one_hot_omics_data(
+        genotype_array=test_array,
+        na_augment_perc=1.0,
+        na_augment_prob=1.0,
+        shuffle_augment_perc=1.0,
+        shuffle_augment_prob=1.0,
+        test_mode=False,
+    )
+
+    assert prepared_array_train != test_array
+    assert (test_array_copy == test_array).all()
+
+    assert (prepared_array_train[:, -1, :] != 1).any()
+
 
 def test_prepare_genotype_array_test_mode():
     test_array = torch.zeros((1, 4, 100), dtype=torch.uint8).detach().numpy()