better tests

jkobject · Sep 5, 2024 · 5c3be00 · 5c3be00
1 parent 152873c
commit 5c3be00
Showing 1 changed file with 38 additions and 38 deletions.
diff --git a/tests/test_base.py b/tests/test_base.py
@@ -13,25 +13,27 @@
 from scdataloader import Preprocessor
 
 from torch.utils.data import DataLoader
+import pytest
 
 
 def test_base():
     assert NAME == "scdataloader"
-
-    print("populating ontology...")
-    start_time = time.time()
-    utils.populate_my_ontology(
-        organisms=["NCBITaxon:10090", "NCBITaxon:9606"],
-        sex=["PATO:0000384", "PATO:0000383"],
-        # celltypes=None,
-        # ethnicities=None,
-        # assays=None,
-        # tissues=None,
-        # diseases=None,
-        # dev_stages=None,
-    )
-    end_time = time.time()
-    print(f"ontology populated in {end_time - start_time:.2f} seconds")
+adata = sc.read_h5ad(os.path.join(os.path.dirname(__file__), "test.h5ad"))
+    try:
+        print("populating ontology...")
+        start_time = time.time()
+        utils.populate_my_ontology(
+            organisms=["NCBITaxon:10090", "NCBITaxon:9606"],
+            sex=["PATO:0000384", "PATO:0000383"],
+            # celltypes=None,
+            # ethnicities=None,
+            # assays=None,
+            # tissues=None,
+            # diseases=None,
+            # dev_stages=None,
+        )
+        end_time = time.time()
+        print(f"ontology populated in {end_time - start_time:.2f} seconds")
     # cx_dataset = (
     #     ln.Collection.using(instance="laminlabs/cellxgene")
     #     .filter(name="cellxgene-census", version="2023-12-15")
@@ -52,26 +54,24 @@ def test_base():
     #     print(i)
     #     break
     # assert True, "Datamodule test passed"
-
-    adata = sc.read_h5ad(os.path.join(os.path.dirname(__file__), "test.h5ad"))
-    preprocessor = Preprocessor(do_postp=False)
-    adata = preprocessor(adata)
-
-    adataset = SimpleAnnDataset(adata, obs_to_output=["organism_ontology_term_id"])
-    col = Collator(
-        organisms=["NCBITaxon:9606"],
-        max_len=1000,
-        how="random expr",
-    )
-    dataloader = DataLoader(
-        adataset,
-        collate_fn=col,
-        batch_size=4,
-        num_workers=1,
-        shuffle=False,
-    )
-    print("dataloader created")
-    for batch in dataloader:
-        print(batch)
-        break
-    assert True, "SimpleAnnDataset test passed"
+        preprocessor = Preprocessor(do_postp=False)
+        adata = preprocessor(adata)
+        adataset = SimpleAnnDataset(adata, obs_to_output=["organism_ontology_term_id"])
+        col = Collator(
+            organisms=["NCBITaxon:9606"],
+            max_len=1000,
+            how="random expr",
+        )
+        dataloader = DataLoader(
+            adataset,
+            collate_fn=col,
+            batch_size=4,
+            num_workers=1,
+            shuffle=False,
+        )
+        print("dataloader created")
+        for batch in dataloader:
+            print(batch)
+            break
+    except Exception as e:
+        pytest.fail(f"An exception occurred: {str(e)}")