update unit tests

ghosttree/scaffold/hybridtree.py

@@ -304,10 +304,6 @@ def _make_foundation_tree(in_name, all_std_error, ghost_tree_fp):
 def _collapse_taxa_line(accession_ids: list, taxonomy: pd.Series,
                         graft_level: int) -> pd.DataFrame:
     """This function was copied and modified from q2-taxa from Qiime2."""
-    if graft_level < 1:
-        raise ValueError('Requested level of %d is too low. Must be greater '
-                         'than or equal to 1.' % graft_level)
-
     # Assemble the taxonomy data
     max_observed_level = _get_max_level(taxonomy)
     if graft_level > max_observed_level:
@@ -318,7 +314,7 @@ def _collapse_taxa_line(accession_ids: list, taxonomy: pd.Series,
     accession_taxonomy_dict = {}
     index_count = 0
     for tax in taxonomy:
-        new_taxa = _collapse(tax, max_observed_level, graft_level)
+        new_taxa = _collapse(tax, graft_level)
         accession_taxonomy_dict[accession_ids[index_count]] = new_taxa
         index_count += 1
 
@@ -330,12 +326,8 @@ def _get_max_level(taxonomy):
     return taxonomy.apply(lambda x: len(x.split(';'))).max()
 
 
-def _collapse(tax, max_observed_level, level):
+def _collapse(tax, level):
     """This function was copied and modified from q2-taxa from Qiime2."""
     tax = [x.strip() for x in tax.split(';')]
-    if len(tax) < max_observed_level:
-        padding = ['__'] * (max_observed_level - len(tax))
-        tax.extend(padding)
-    # return ';'.join(tax[:level])
     taxa = tax[:level][-1].split('__')[1].capitalize()
     return taxa

ghosttree/scaffold/tests/test_hybridtree.py

@@ -24,6 +24,7 @@ def setUp(self):
         self.extension_taxonomy = StringIO(extension_taxonomy)
         self.extension_taxonomy_unids = StringIO(extension_taxonomy_unids)
         self.extension_taxonomy_none = StringIO(extension_taxonomy_none)
+        self.extension_tax_lots_unids = StringIO(extension_tax_lots_unids)
         self.extension_seqs = StringIO(extension_seqs)
         self.foundation_alignment = StringIO(foundation_alignment)
         self.ghost_tree_fp = StringIO(ghost_tree_fp)
@@ -109,6 +110,14 @@ def test_extension_genus_accession_dic_unidentifieds_level_6_genus(self):
                                                  self.graft_level_6)
         self.assertDictEqual(result, test)
 
+    def test_extension_genus_accession_dic_lots_unids_level_6_genus(self):
+        test = {'Mucor': ['C1', 'C2', 'C3', 'M1', 'M2', 'M3', 'C4', 'C5',
+                          'M4'], 'Phoma': ['P1', 'P2']}
+        result = _extension_genus_accession_dict(self.otu_clusters,
+                                                 self.extension_tax_lots_unids,
+                                                 self.graft_level_6)
+        self.assertDictEqual(result, test)
+
     def test_extension_genus_accession_dic_unidentifieds_level_5_family(self):
         # There is a tie breaker here so data can switch to either
         # Saccharomycetaceae OR Mucoraceae randomly (no priority).
@@ -169,6 +178,19 @@ def test_extension_genus_accession_dic_unidentifieds_level_5_family(self):
 M4\tk__Fungi;p__As;c__Do;o__My;f__Mucoraceae;g__Mucor;s__El
 """
 
+extension_tax_lots_unids = """P1\tk__Fungi;p__As;c__Do;o__My;f__Didymellaceae;g__Phoma;s__El
+P2\tk__Fungi;p__As;c__Do;o__My;f__Didymellaceae;g__Phoma;s__El
+C1\tk__Fungi;p__As;c__Do;o__My;f__Saccharomycetaceae;g__Unidentified;s__El
+C2\tk__Fungi;p__As;c__Do;o__My;f__Saccharomycetaceae;g__Unidentified;s__El
+C3\tk__Fungi;p__As;c__Do;o__My;f__Saccharomycetaceae;g__Unidentified;s__El
+C4\tk__Fungi;p__As;c__Do;o__My;f__Saccharomycetaceae;g__Unidentified;s__El
+C5\tk__Fungi;p__As;c__Do;o__My;f__Unidentified;g__Unidentified;s__El
+M1\tk__Fungi;p__As;c__Do;o__My;f__Mucoraceae;g__Mucor;s__El
+M2\tk__Fungi;p__As;c__Do;o__My;f__Mucoraceae;g__Mucor;s__El
+M3\tk__Fungi;p__As;c__Do;o__My;f__Mucoraceae;g__Mucor;s__El
+M4\tk__Fungi;p__As;c__Do;o__My;f__Mucoraceae;g__Mucor;s__El
+"""
+
 extension_taxonomy_none = """P1\tk__Fungi;p__As;c__Do;o__My;f__Els;Phoma;s__El
 P2\tk__Fungi;p__As;c__Do;o__My;f__Els;_phoma;s__El
 """