Reasonable tests for alignments module

Author: jeromekelleher

sc2ts/alignments.py

@@ -80,7 +80,7 @@ class AlignmentStore(collections.abc.Mapping):
     def __init__(self, path, mode="r"):
         map_size = 1024**4
         self.env = lmdb.Environment(
-            path, subdir=False, readonly=mode == "r", map_size=map_size
+            str(path), subdir=False, readonly=mode == "r", map_size=map_size
         )
 
     def __enter__(self):
@@ -95,21 +95,6 @@ def close(self):
     def __str__(self):
         return f"AlignmentStore at {self.env.path()} contains {len(self)} alignments"
 
-    @staticmethod
-    def initialise(path):
-        """
-        Create a new store at this path.
-        """
-        db_path = pathlib.Path(path)
-        if db_path.exists():
-            db_path.unlink()
-
-        reference = core.get_reference_sequence()
-        with lmdb.Environment(str(db_path), subdir=False) as env:
-            with env.begin(write=True) as txn:
-                txn.put("MN908947".encode(), compress_alignment(reference))
-        return AlignmentStore(path, "a")
-
     def _flush(self, chunk):
         logger.debug(f"Flushing {len(chunk)} sequences")
         with self.env.begin(write=True) as txn:
@@ -157,20 +142,6 @@ def __len__(self):
         with self.env.begin() as txn:
             return txn.stat()["entries"]
 
-    def get_all(self, strains, sequence_length):
-        A = np.zeros((len(strains), sequence_length), dtype=np.int8)
-        with self.env.begin() as txn:
-            for j, strain in enumerate(strains):
-                val = txn.get(strain.encode())
-                if val is None:
-                    raise KeyError(f"{strain} not found")
-                a = decompress_alignment(val)
-                if len(a) != sequence_length:
-                    raise ValueError(
-                        f"Alignment for {strain} not of length {sequence_length}"
-                    )
-        return A
-
 
 @dataclasses.dataclass
 class MaskedAlignment:

tests/test_alignments.py

@@ -1,12 +1,67 @@
+import pathlib
+import shutil
+import gzip
+
 import numpy as np
 import pytest
 from numpy.testing import assert_array_equal
 
-# FIXME - quick hacks here to get tests working
 from sc2ts import alignments as sa
 from sc2ts import core
 
 
+@pytest.fixture
+def data_cache():
+    cache_path = pathlib.Path("tests/data/cache")
+    if not cache_path.exists():
+        cache_path.mkdir()
+    return cache_path
+
+
+@pytest.fixture
+def alignments_fasta(data_cache):
+    cache_path = data_cache / "alignments.fasta"
+    if not cache_path.exists():
+        with gzip.open("tests/data/alignments.fasta.gz") as src:
+            with open(cache_path, "wb") as dest:
+                shutil.copyfileobj(src, dest)
+    return cache_path
+
+
+@pytest.fixture
+def alignments_store(data_cache, alignments_fasta):
+    cache_path = data_cache / "alignments.db"
+    if not cache_path.exists():
+        with sa.AlignmentStore(cache_path, "a") as a:
+            fasta = core.FastaReader(alignments_fasta)
+            a.append(fasta, show_progress=False)
+    return sa.AlignmentStore(cache_path)
+
+
+class TestAlignmentsStore:
+    def test_info(self, alignments_store):
+        assert "contains" in str(alignments_store)
+
+    def test_len(self, alignments_store):
+        assert len(alignments_store) == 55
+
+    def test_fetch_known(self, alignments_store):
+        a = alignments_store["SRR11772659"]
+        assert a.shape == (core.REFERENCE_SEQUENCE_LENGTH,)
+        assert a[0] == "X"
+        assert a[1] == "N"
+        assert a[-1] == "N"
+
+    def test_keys(self, alignments_store):
+        keys = list(alignments_store.keys())
+        assert len(keys) == len(alignments_store)
+        assert "SRR11772659" in keys
+
+    def test_in(self, alignments_store):
+        assert "SRR11772659" in alignments_store
+        assert "NOT_IN_STORE" not in alignments_store
+
+
 def test_get_gene_coordinates():
     d = core.get_gene_coordinates()
     assert len(d) == 11
@@ -66,6 +121,12 @@ def test_examples(self, a):
         with pytest.raises(ValueError):
             sa.decode_alignment(np.array(a))
 
+    def test_encode_real(self, alignments_store):
+        h = alignments_store["SRR11772659"]
+        a = sa.encode_alignment(h)
+        assert a[0] == -1
+        assert a[-1] == -1
+
 
 class TestMasking:
     # Window size of 1 is weird because we have to have two or more
@@ -113,3 +174,20 @@ def test_bad_window_size(self, w):
         a = np.zeros(2, dtype=np.int8)
         with pytest.raises(ValueError):
             sa.mask_alignment(a, window_size=w)
+
+
+class TestEncodeAndMask:
+    def test_known(self, alignments_store):
+        a = alignments_store["SRR11772659"]
+        ma = sa.encode_and_mask(a)
+        assert ma.original_base_composition == {
+            "T": 9566,
+            "A": 8894,
+            "G": 5850,
+            "C": 5472,
+            "N": 121,
+        }
+        assert ma.original_md5 == "e96feaa72c4f4baba73c2e147ede7502"
+        assert len(ma.masked_sites) == 133
+        assert ma.masked_sites[0] == 1
+        assert ma.masked_sites[-1] == 29903