add unit tests

ghosttree/silva/filter.py

@@ -13,7 +13,7 @@ def fungi_from_fasta(fasta_fh, accession_fh, taxonomy_fh):
 def _parse_accession_map(accession_fh):
     accession_map = {}
     for line in accession_fh:
-        accession, map_num = line.split("\t")
+        accession, map_num = line.strip().split("\t")
         if accession in accession_map:
             raise ValueError("Duplicate accession number %r" % accession)
         accession_map[accession] = map_num
@@ -23,7 +23,7 @@ def _parse_accession_map(accession_fh):
 def _parse_taxonomy_map(taxonomy_fh):
     taxonomy_map = {}
     for line in taxonomy_fh:
-        split_line = line.split("\t")
+        split_line = line.strip().split("\t")
         if len(split_line) == 3:
             taxonomy, map_num, rank = split_line
             if map_num in taxonomy_map:

ghosttree/silva/tests/test_filter.py

@@ -1,23 +1,99 @@
-
 import unittest
 from StringIO import StringIO
 
+from skbio import BiologicalSequence
+
 from ghosttree.silva.filter import fungi_from_fasta
 
 
 class TestFungiFromFasta(unittest.TestCase):
+    def setUp(self):
+        self.fasta_no_fungi = StringIO(fasta_no_fungi)
+        self.fasta_with_fungi = StringIO(fasta_with_fungi)
+        self.fasta_many_fungi = StringIO(fasta_many_fungi)
+        self.accession = StringIO(accession)
+        self.accession_with_duplicates = StringIO(accession_with_duplicates)
+        self.taxonomy_no_fungi = StringIO(taxonomy_no_fungi)
+        self.taxonomy_with_fungi = StringIO(taxonomy_with_fungi)
+        self.taxonomy_with_duplicates = StringIO(taxonomy_with_duplicates)
+
     def test_fasta_no_fungi(self):
-        fasta_file = StringIO(fasta_no_fungi)
-        accession_file = StringIO(accession)
-        taxonomy_file = StringIO(taxonomy)
-        result = fungi_from_fasta(fasta_file, accession_file, taxonomy_file)
+        result = fungi_from_fasta(self.fasta_no_fungi,
+                                  self.accession,
+                                  self.taxonomy_no_fungi)
         self.assertEqual(list(result), [])
 
+    def test_fasta_with_fungi(self):
+        result = fungi_from_fasta(self.fasta_with_fungi,
+                                  self.accession,
+                                  self.taxonomy_with_fungi)
+        self.assertEqual(list(result),
+                         [BiologicalSequence("ATCG", id="AB21",
+                                             description="Fungi")])
+
+    def test_fasta_with_many_fungi(self):
+        result = fungi_from_fasta(self.fasta_many_fungi,
+                                  self.accession,
+                                  self.taxonomy_with_fungi)
+        self.assertEqual(list(result), [
+            BiologicalSequence("GGGG", id="AB123", description="Fungi"),
+            BiologicalSequence("CCCC", id="AB125", description="Fungi"),
+            BiologicalSequence("AAAA", id="AB126", description="Fungi"),
+        ])
+
+    def test_duplicate_accession_numbers(self):
+        with self.assertRaises(ValueError):
+            list(fungi_from_fasta(self.fasta_with_fungi,
+                                  self.accession_with_duplicates,
+                                  self.taxonomy_with_fungi))
+
+    def test_duplicate_map_numbers(self):
+        with self.assertRaises(ValueError):
+            list(fungi_from_fasta(self.fasta_with_fungi,
+                                  self.accession,
+                                  self.taxonomy_with_duplicates))
+
+
+fasta_many_fungi = """>AB123 Fungi
+GGGG
+>AB124 Bacteria
+TTTT
+>AB125 Fungi
+CCCC
+>AB126 Fungi
+AAAA
+>AB100 Fungi
+ATAC
+"""
+
 fasta_no_fungi = """>AB21 Bacteria
 ATCG
 """
-accession = "AB21\t123"
-taxonomy = "Bacteria\t123\tdomain"
+fasta_with_fungi = """>AB21 Fungi
+ATCG
+"""
+taxonomy_with_fungi = """Fungi\t123\tgenus
+Eukaryote;Fungi;Fungal Species\t456\tgenus
+Fungi;Fungal species\t789\tgenus
+Bacteria;bacterial species\t111\tgenus
+Fungi;cladosporium\t100\tdomain
+"""
+taxonomy_no_fungi = "Bacteria\t123\tdomain"
+taxonomy_with_duplicates = """Fungi\t123\tgenus
+Fungi\t123\tgenus
+Fungi\t122\tgenus
+"""
+accession = """AB21\t123
+AB123\t456
+AB124\t111
+AB125\t789
+AB126\t123
+AB100\t100
+"""
+accession_with_duplicates = """AB21\t123
+AC33\t456
+AC33\t789
+"""
 
 
 if __name__ == "__main__":