Add tests

q2_moshpit/__init__.py

@@ -12,7 +12,7 @@
 from . import prodigal
 from ._version import get_versions
 from .dereplication import dereplicate_mags
-from .filter import filter_derep_mags, filter_mags
+from .filtering import filter_derep_mags, filter_mags
 from .kaiju import classification as kaiju_class, database as kaiju_db
 from .kraken2 import (
     classification as kraken_class,

q2_moshpit/filter/filter_mags.py → q2_moshpit/filtering/filter_mags.py

@@ -36,15 +36,14 @@ def _filter_ids(
     Returns:
         set: The filtered set of IDs.
     """
-    if metadata:
-        selected_ids = metadata.get_ids(where=where)
-        if not selected_ids:
-            print("The filter query returned no IDs to filter out.")
+    selected_ids = metadata.get_ids(where=where)
+    if not selected_ids:
+        print("The filter query returned no IDs to filter out.")
+    else:
+        if exclude_ids:
+            ids -= set(selected_ids)
         else:
-            if exclude_ids:
-                ids -= set(selected_ids)
-            else:
-                ids &= set(selected_ids)
+            ids &= set(selected_ids)
     print(f"Found {len(ids)} IDs to keep.")
     return ids
 
@@ -75,7 +74,6 @@ def _filter_manifest(
         manifest.index.get_level_values('sample-id') + "/" + \
         manifest.index.get_level_values('mag-id') + ".fasta"
 
-
     return manifest[manifest.index.get_level_values(lvl).isin(ids_to_keep)]
 
 
@@ -147,11 +145,14 @@ def filter_mags(
     )
     try:
         for _id, row in filtered_mags.iterrows():
-            sample_dir = os.path.join(str(results), row["sample_id"])
+            if on == 'mag':
+                sample_dir = os.path.join(str(results), row["sample_id"])
+                mag_dest = os.path.join(sample_dir, f"{_id}.fasta")
+            else:
+                sample_dir = os.path.join(str(results), _id)
+                mag_dest = os.path.join(sample_dir, f"{row['mag_id']}.fasta")
             os.makedirs(sample_dir, exist_ok=True)
-            duplicate(
-                row["mag_fp"], os.path.join(sample_dir, f"{_id}.fasta")
-            )
+            duplicate(row['mag_fp'], mag_dest)
     except KeyError:
         raise ValueError(f"{_id!r} is not a MAG present in the input data.")
 

q2_moshpit/filtering/tests/__init__.py

@@ -0,0 +1,7 @@
+# ----------------------------------------------------------------------------
+# Copyright (c) 2022-2023, QIIME 2 development team.
+#
+# Distributed under the terms of the Modified BSD License.
+#
+# The full license is in the file LICENSE, distributed with this software.
+# ----------------------------------------------------------------------------

q2_moshpit/filtering/tests/data/MANIFEST

@@ -0,0 +1,7 @@
+sample-id,mag-id,filename
+id1,mag1,id1/mag1.fasta
+id1,mag2,id1/mag2.fasta
+id2,mag3,id2/mag3.fasta
+id3,mag4,id3/mag4.fasta
+id3,mag5,id3/mag5.fasta
+id3,mag6,id3/mag6.fasta

q2_moshpit/filtering/tests/data/mags/MANIFEST

@@ -0,0 +1,4 @@
+sample-id,mag-id,filename
+sample1,24dee6fe-9b84-45bb-8145-de7b092533a1,sample1/24dee6fe-9b84-45bb-8145-de7b092533a1.fasta
+sample2,db03f8b6-28e1-48c5-a47c-9c65f38f7357,sample2/db03f8b6-28e1-48c5-a47c-9c65f38f7357.fasta
+sample2,d65a71fa-4279-4588-b937-0747ed5d604d,sample2/d65a71fa-4279-4588-b937-0747ed5d604d.fasta

q2_moshpit/filtering/tests/data/mags/sample1/24dee6fe-9b84-45bb-8145-de7b092533a1.fasta

@@ -0,0 +1,4 @@
+>NZ_00000000.1_contig1
+ATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATGTTTCCAGATGCAATGCGTGGGCACTGATG
+>NZ_00000000.1_contig2
+TTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGAAAACCATGCAGTGTTGACGTCAGTGA

q2_moshpit/filtering/tests/data/mags/sample2/d65a71fa-4279-4588-b937-0747ed5d604d.fasta

@@ -0,0 +1,6 @@
+>NZ_CP018863.1_contig1
+GCCTCCTCCCAGTTCGTCTCAGCGCTGCTGCTGGTCGGCGCCAAATTCCGTAACGGGCTGCACCTCGAACATTCCGGCCAGAGCGTCCCCAGCCTGGACCACGTTGCCATGACCGTGGCGGTACTGCGCAGCGTCGGGGTGGAGGTAGACGATTCCCGGCAGAACCACTGGTGGTCCGGCCCGGACCGGTCAAGGCCTTCGACGTCACCGTCGAACAGGACCTTTCCAATGCCGGCCCCTTCCTCGCGGCAGCCCTGGCCACGAAGGGAACGGTTCGGATCCCAGGCTGGCCGGAGAAAACCACGCAGGTAGGTGACAAATGGCGCAGCATCCTGGCGCAACTCGGCGCCACTGTCAGCTACGAGAACGGCACCCTCACCGTAACCGGCGGGGCAGAGATCACCGGGGCGCAGCTCGCCGACACCAGCGAACTTGCCCCCACCACGGCGGCGCTCTGTGCCCTGGCCGGCAGCGAATCCAGGCTCACCGGAATTGCCCACTTGCGGGGACACGAAACAGACCGGCTGGCGGCTCTGGTCGCGGAAATCAATGCCTTGGGTGGCGACGCCGAAGAAACCGAAGACGGGTTGATCATCCGTCCGAGGCCACTGCATGGCGGGGTCTTCCATTCATATGAGGACCACAGGATGGCCACCGCCGGAGCCATTATCGGGCTGGCAGTCGAGGGCGTGGAAGTCGAAGACATCGGCACCACGGCCAAGACCATGCCCGAGTTCCCGCGGCTATGGCAGGACCTGTTCGAGACTTCCGTCCGCCAGTCCGAGGCGGGAGCGCTCTAAGGTGGTGCGCGGCAACCGTACGTGGGACGAGTCCGATGTCCGCATCCGTCCCAACAAGCGCGGCTCGCGTCCGCGTACCAAGGAACGGCCTGCCCACGAAGACGCCGTCATCGGGCGGATCATCACCGTGGACCGCGGCCGCTACACCGCGGTCGTCGATGAAGACACTGCCCGGGAACGGGTGGTTGTCGCCGCCCGTG
+>NZ_CP018863.1_contig2
+CCCGGGAGCTTCGTCGCAGTCCGGTGGTGGCCGGCGACTTCGTAGCGCTCGTCGGTGACATTACCGGTGAGCCGGATACGTTGGCCCGGCTGGTCCGGATTGAGGAACGCCGGACGCTGTTGCGCCGCAGCGCCGACGATACAGATCCCGTGGAGCGGGTAGTCGTCGCCAATGCAGACCAGCTGGTCATCGTCGTGGCCGCCGCAAACCCCGAGCCGCGCACCGGTTTCATCGACCGCGCCTTGGTAGCGGCGTACGACGCCGGTATCAGCCCGCTGCTGTGCGTCACCAAAGCGGACGTCAAGGATCCCGAAGAACTGCTCTCCAACTACCGGCACCTGGACCTGCCCGTGATCGTCAGCCGGACGGCCGGCACGGAGGGCTCCGGGGTGGATGCACGGTCCGCCGACGGGCTGTCTGCCCGTCTCGACCGTGACGCCGTAGCGGCGCTCCGTGGCTATCTGGATGGGATGGTCAGCGTCATGCTCGGCCATTCGGGCGTGGGCAAGTCCACCATGGTCAATGCCCTCACGGGGGCGGAGCGCGCCACGGGGGGAGTCAACGCGGTGACCGGGCGGGGCCGGCATACCTCCTCCTCGGCGCTGGCCCTGAAGCTGGCCGACGCTCCGGCTGGCAGCTGGATCATCGACACGCCCGGCATCCGTTCTTTTGGACTGGCCCACGTGGACCCGGACCGGATCATTTCCGCTTTTCCCGATTTGGAGCCCGGGACGGCGGACTGCGAGCGGGGCTGCAAGCACGACGACCATGCCGTCAACTGCGGCGTGGACGCCTGGGTGGCCTCCGGGCAGGCCGGCGAATCCGGCCCGGCACGGCTGGCCTCGCTGCGCCGTTTGCTGGGAACGGAAGAACGCGCCCAGGCGAAGGAACTCGGGTTCCAGTAGCACCGCCGTCGTCGGTCAGGGACTTCACATCCCGCATCCGGCCGCCAAATAAGGATAAGTTGAAGCCTATGACCCGTGACGTTCAAAGCTATAAC
+>NZ_CP018863.1_contig3
+GACGATCTGCGCCTGGCCCATGTGATGGCCGATTCCGTGGATTCGCAGACCATGGCCCGCTTCAAGGCGCTGGACCTGAAAATCGAGACCAAGCCGGATCTCACCCCTGTCACGGATGCGGACCGCGCCGCTGAAGAGGCCATCCGCGGCCAACTCTCCCGGGCCCGGCCGCGCGACGCGGTCCTCGGCGAGGAATACGGCAGCAGCGGCCACGGCTCCCGCCGCTGGATCATCGATCCCATCGACGGCACGAAGAACTTCGTCCGCGGGGTGCCGGTCTGGGCCACCTTGATCGCGCTGGTAGACGAAGACCGTCCCGTGGTCGGCCTGGTCAGCGCGCCGGCTCTGGGCAAGCGCTGGTGGGCCGCGACCGGAACCGGTGCCTACATGGGACGTTCGCTGTCCGCGGCCACCCGGCTCCGGGTATCCGATGTCAACCGGCTCGAGGACGCGTCCCTCTCCTATTCCAGCCTCACCGGCTGGCAGGAACGCGGCAACTTCCCGGAGTTCCTCGGCCTCACCGAATCCGTCTGGCGCACCCGTGCCTACGGGGACTTCTGGTCCTACTGCATGGTGGCCGAGGGCGCCGTCGACATTGCCTGCGAACCCGAACTCAACCTCTATGACATGGCGGCCCTCGTGCCGATCGTGACCGAGGCCGGCGGACGGTTCAGTTCGCTCGAGGGCGAGGACGGACCCTTCGGCGGCAACGCGTTGGCCACGAACGGCACGCTGCACGACGAGGTCCTCTACCGGCTCAATCCGCAGTTGCGCGGCCAGCGTCCGGCCGCACACCCGGAGGACGGGTCCCTGCCGGAAACCGCTCCGGAGGCCTCCATGGAGGCGGACGGCCTGCGCTGACGCTGTCTTTTGTGACGAATTACGACGGCGGCCGTCCCCATTCCGGGGATGGCCGCCTTTTCGTTCCCGTAACAAAGATGCGGCCCCTCCGGCCGGACAATAATCTCGATGGCAGGTCACGAGTGCCAGCGCTAAACCC

q2_moshpit/filtering/tests/data/mags/sample2/db03f8b6-28e1-48c5-a47c-9c65f38f7357.fasta

@@ -0,0 +1,4 @@
+>NZ_CP007255.1_contig1
+GGCGTGCCCTCAGCGGTCTCCTGCCCATTTACCCGGTCCATCAGACCGCGGACATACTCTTCTGCAAAGGCAACCTCGTTCCAGTCGGGAAAGACAACCTCCCACACGTCGAAATGACCCGTGTCATCGCGCGACGCTTCAACGACCGCTACGGCCAGGTCTTCCCCGAACCTCAAGCGCTGCTTACCTCGAGCCCCGAAGTACCTGAACTCGACGGCAGAAAGATATCCAAGAGCTACGGAAACGCCATCACACTGTCGATGACCGCCGACGACACCGCAGCGGCAATCCGCAAGGCACCGACCGACTCACAACGCCGCATCACCTTCGACCCGATCAACCGCCCCGGGGTATCCGCACTGCTGACCACAGCAGCATTGTGCTCAGGTCGCGACGAGAATGACCTCGCCGACCAGATCGGCGACGCCGGCGCAGGCGCACTCAAAAAGCTCACCACCACGCTCGGGAACGACTTCCTCGCCCCTCACCGGCAACGCAGAAACGACCTCGCCCAAGACCCCCACCATGTCGCCGACGTCCTGCGCATCGGAAACCAGCGCGCGAACGAAATCGCCCAACAGACCCTCGAGGAAGTCCGAACCGCAATGGGCACCATCTACTGATCGACTCCTACCCATAGCGCACTCCAGTAGTTTGGACATTCACCGGGTTGCGACGATCCCTGGAATCCGCCATCGGGAGTGTCTGATCAACGGTGGATCCGACCTCGGTATCGGTTAATTACTGTTGAGGGTGATGCGTGCTTCGAAGGTGATAGCGAACGCGTTCAGCGCTGGCTTCCACCTCATCGCCCATCGCGCCTTACCTTTCCCAGTCGGGTCCAGAGCCCGAGTCGCCAGATACAAGCATTCGAGCGCCGCCTGCTCCGTCGGGAAGTGCCCACGAGCACGCACAGCACGCCGGTAACGAGCGTTGACGGACTCGATCGCATTGGTCGAGCAGATCACCCGCAGGATCTCTTTCGTCGTAATCCAAGAAC
+>NZ_CP007255.1_contig3
+ATGATCGATCCTTCCCACCAAGCTCAGCTCGGCGTGTCGAACCAAGACCGGATCAACCGTTTATCTGACAGTCCCCTCACTAGCGAGTAAGTCGTCAGAGCGATTCTTCGGGTTTGGGTGTGCGTTTTTGACCATCGCCATACCGGACAGCACGGAACGGGCTTCTGGGCGTGACTGACGTGGAGAAACTTAGCGAAAAGTTTCCAATTGGGTGTCCCGTCGGCGAACGTGGGTGTGAAGATCGGGTGTCGAGATGCAACACCGCCGGCAGAACACGGGATCCAGGTTGATCTTCTGGTGAAGGCGGGTGTGGATCGCGACCAGGCTTATGTCGATCATCGTGTCAGCGGCGGGCAGGCTCGTCGGCCAGGTTTGGATGATGCGATCAACGCGGCTCGGGCCGGCGACGTCCTTTGTGTGACGATGCCTCATTAACTGTCGCGCTCGGCGAAGGATCTGCATGAGCCGGTAGAACGGATCGCGGAGAAGGTGTATCGCTGTCAATTGACGGTCAGTTGTACGACCCGCGCGATCCGATGGGGAAGATGTTCATCGGCTTGCTGGGGATTGATGACGGAGTTTCGAGTCGGATCTGATTCGTGGCCGCACTCGTGACGCCCTGGCCGCCGCGACAGCTGCTGAGAAGAGAAAGGTGCGTCGGACAAGCTCACCGGCTGGAGAAAAAAGTCAGTCACGATCGGGTCTTTTCTTCCAGAGGCTCCCGGTCTGGAGCCTCACCGGTTGCCGCGGGTCGGGCTCCTTCGATGTGTGGGTCAGCTCGTTACGCCGCGTTAAGAATCTGAAAACTTGGTGAGCAAACCTTCCTTTCCCAGGTGAAGCGGCGGATTATCGAAGTGGCCGCGCATCTGGTTACCGCGCGGCATGTCGAGCGCGCCTGAAGGCACAACGGGCCGGGAGGTCACATGCTGCTACACCAATCGATAGGGTTGGAGCGCTTCAACGAACTACCGCGACAGAAGGCGGTTCACGCACTGTTCGA

q2_moshpit/filtering/tests/data/metadata-derep.tsv

@@ -0,0 +1,4 @@
+mag_id	sample_id	input_file	dataset	complete	single	duplicated	fragmented	missing	n_markers	scaffold_n50	contigs_n50	percent_gaps	scaffolds	length
+24dee6fe-9b84-45bb-8145-de7b092533a1	sample1	24dee6fe-9b84-45bb-8145-de7b092533a1.fasta	bacteria_odb10	28.2	27.4	0.8	8.9	62.9	124	4785	4785	0.000%	265	1219165
+d65a71fa-4279-4588-b937-0747ed5d604d	sample2	d65a71fa-4279-4588-b937-0747ed5d604d.fasta	bacteria_odb10	1.6	1.6	0.0	1.6	96.8	124	3548	3548	0.000%	67	245922
+db03f8b6-28e1-48c5-a47c-9c65f38f7357	sample2	db03f8b6-28e1-48c5-a47c-9c65f38f7357.fasta	bacteria_odb10	26.6	26.6	0.0	3.2	70.2	124	78679	78679	0.000%	17	714893

q2_moshpit/filtering/tests/data/metadata-sample.tsv

@@ -0,0 +1,3 @@
+id	input_file	dataset	metric
+sample1	24dee6fe-9b84-45bb-8145-de7b092533a1.fasta	bacteria_odb10	28.2
+sample2	d65a71fa-4279-4588-b937-0747ed5d604d.fasta	bacteria_odb10	1.6

q2_moshpit/filtering/tests/test_filter.py

@@ -0,0 +1,170 @@
+# ----------------------------------------------------------------------------
+# Copyright (c) 2022-2023, QIIME 2 development team.
+#
+# Distributed under the terms of the Modified BSD License.
+#
+# The full license is in the file LICENSE, distributed with this software.
+# ----------------------------------------------------------------------------
+import unittest
+
+import pandas as pd
+import qiime2
+from qiime2.plugin.testing import TestPluginBase
+
+from q2_moshpit.busco.types import BUSCOResultsFormat
+from q2_moshpit.filtering.filter_mags import (
+    _filter_ids, _filter_manifest, _mags_to_df,
+    filter_derep_mags, filter_mags
+)
+from q2_types.feature_data_mag import MAGSequencesDirFmt
+from q2_types.per_sample_sequences import MultiMAGSequencesDirFmt
+
+
+class TestMAGFiltering(TestPluginBase):
+    package = 'q2_moshpit.filtering.tests'
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.metadata_df = pd.DataFrame({
+            'sample_id': ['id1', 'id1', 'id2', 'id3', 'id3', 'id3'],
+            'complete': [20.0, 98.0, 68.5, 100.0, 98.0, 21.3]
+        }, index=pd.Index(
+            ['mag1', 'mag2', 'mag3', 'mag4', 'mag5', 'mag6'],
+            name='id'
+        ))
+
+    def setUp(self):
+        super().setUp()
+        self.manifest_df = pd.read_csv(self.get_data_path("MANIFEST"))
+        self.manifest_df.set_index(['sample-id', 'mag-id'], inplace=True)
+        self.mag_data_dir = self.get_data_path('mags')
+        self.mag_derep_data_dir = self.get_data_path('mags/sample2')
+        self.mag_df = pd.DataFrame({
+            'sample_id': ['sample1', 'sample2', 'sample2'],
+            'mag_id': [
+                '24dee6fe-9b84-45bb-8145-de7b092533a1',
+                'd65a71fa-4279-4588-b937-0747ed5d604d',
+                'db03f8b6-28e1-48c5-a47c-9c65f38f7357'
+            ],
+            'mag_fp': [
+                f'{self.mag_data_dir}/sample1/'
+                f'24dee6fe-9b84-45bb-8145-de7b092533a1.fasta',
+                f'{self.mag_data_dir}/sample2/'
+                f'd65a71fa-4279-4588-b937-0747ed5d604d.fasta',
+                f'{self.mag_data_dir}/sample2/'
+                f'db03f8b6-28e1-48c5-a47c-9c65f38f7357.fasta'
+            ]
+        })
+        self.metadata = qiime2.Metadata(self.metadata_df)
+
+    def test_filter_ids_include(self):
+        ids = {'mag1', 'mag2', 'mag3', 'mag4', 'mag5', 'mag6'}
+        obs = _filter_ids(
+            ids, self.metadata, 'complete>40', exclude_ids=False
+        )
+        exp = {'mag2', 'mag3', 'mag4', 'mag5'}
+        self.assertSetEqual(obs, exp)
+
+    def test_filter_ids_include_none(self):
+        ids = {'mag1', 'mag2', 'mag3', 'mag4', 'mag5', 'mag6'}
+        obs = _filter_ids(
+            ids, self.metadata, 'complete<10', exclude_ids=False
+        )
+        exp = ids
+        self.assertSetEqual(obs, exp)
+
+    def test_filter_ids_exclude(self):
+        ids = {'mag1', 'mag2', 'mag3', 'mag4', 'mag5', 'mag6'}
+        obs = _filter_ids(
+            ids, self.metadata, 'complete>40', exclude_ids=True
+        )
+        exp = {'mag1', 'mag6'}
+        self.assertSetEqual(obs, exp)
+
+    def test_filter_manifest_mags(self):
+        ids = {'mag1', 'mag2', 'mag5'}
+        obs = _filter_manifest(self.manifest_df, ids, on='mag')
+        exp = self.manifest_df[
+            self.manifest_df.index.get_level_values('mag-id').isin(ids)
+        ]
+        pd.testing.assert_frame_equal(obs, exp)
+
+    def test_filter_manifest_samples(self):
+        ids = {'id1', }
+        obs = _filter_manifest(self.manifest_df, ids, on='sample')
+        exp = self.manifest_df[
+            self.manifest_df.index.get_level_values('sample-id').isin(ids)
+        ]
+        pd.testing.assert_frame_equal(obs, exp)
+
+    def test_filter_manifest_error(self):
+        ids = {'id1', }
+        with self.assertRaisesRegex(ValueError, 'parameter: unicorn'):
+            _filter_manifest(self.manifest_df, ids, on='unicorn')
+
+    def test_mags_to_df_on_sample(self):
+        mags = MultiMAGSequencesDirFmt(self.mag_data_dir, mode='r')
+        obs = _mags_to_df(mags, on='sample')
+        exp = self.mag_df
+        exp.set_index('sample_id', inplace=True)
+        pd.testing.assert_frame_equal(obs, exp)
+
+    def test_mags_to_df_on_mag(self):
+        mags = MultiMAGSequencesDirFmt(self.mag_data_dir, mode='r')
+        obs = _mags_to_df(mags, on='mag')
+        exp = self.mag_df
+        exp.set_index('mag_id', inplace=True)
+        pd.testing.assert_frame_equal(obs, exp)
+
+    def test_filter_derep_mags(self):
+        mags = MAGSequencesDirFmt(self.mag_derep_data_dir, mode='r')
+        metadata = BUSCOResultsFormat(
+            self.get_data_path('metadata-derep.tsv'), mode='r'
+        ).view(qiime2.Metadata)
+
+        obs = filter_derep_mags(mags, metadata, where='complete>10')
+        obs_features = obs.feature_dict()
+        exp_features = ['db03f8b6-28e1-48c5-a47c-9c65f38f7357']
+        self.assertListEqual(list(obs_features.keys()), exp_features)
+
+    def test_filter_mags_features(self):
+        mags = MultiMAGSequencesDirFmt(self.mag_data_dir, mode='r')
+        metadata = BUSCOResultsFormat(
+            self.get_data_path('metadata-derep.tsv'), mode='r'
+        ).view(qiime2.Metadata)
+
+        obs = filter_mags(mags, metadata, where='length<1000000', on='mag')
+        obs_samples = obs.sample_dict()
+        exp_samples = ['sample2']
+        self.assertListEqual(list(obs_samples.keys()), exp_samples)
+
+        obs_features = [
+            mag_id for feature_dict in obs_samples.values()
+            for mag_id in feature_dict.keys()
+        ]
+        exp_features = [
+            'd65a71fa-4279-4588-b937-0747ed5d604d',
+            'db03f8b6-28e1-48c5-a47c-9c65f38f7357'
+        ]
+        self.assertListEqual(obs_features, exp_features)
+
+    def test_filter_mags_samples(self):
+        mags = MultiMAGSequencesDirFmt(self.mag_data_dir, mode='r')
+        metadata = qiime2.Metadata(
+            pd.read_csv(self.get_data_path('metadata-sample.tsv'),
+                        sep='\t', index_col=0)
+        )
+
+        obs = filter_mags(mags, metadata, where='metric<5', on='sample')
+        obs_samples = obs.sample_dict()
+        exp_samples = ['sample2']
+        self.assertListEqual(list(obs_samples.keys()), exp_samples)
+
+        obs_feature_count = len(obs.sample_dict()['sample2'])
+        exp_feature_count = 2
+        self.assertEqual(obs_feature_count, exp_feature_count)
+
+
+if __name__ == "__main__":
+    unittest.main()

q2_moshpit/plugin_setup.py

@@ -1216,7 +1216,7 @@
                    "excluded from the filtered data.",
 }
 plugin.methods.register_function(
-    function=q2_moshpit.filter.filter_derep_mags,
+    function=q2_moshpit.filtering.filter_derep_mags,
     inputs={"mags": FeatureData[MAG]},
     parameters=filter_mags_params,
     outputs={"filtered_mags": FeatureData[MAG]},
@@ -1227,7 +1227,7 @@
 )
 
 plugin.methods.register_function(
-    function=q2_moshpit.filter.filter_mags,
+    function=q2_moshpit.filtering.filter_mags,
     inputs={"mags": SampleData[MAGs]},
     parameters={
         **filter_mags_params,

setup.py

@@ -87,7 +87,10 @@
         'q2_moshpit.partition.tests': [
             'data/*', 'data/*/*', 'data/*/*/*', 'data/*/*/*/*',
             'data/*/*/*/*/*',
-        ]
+        ],
+        'q2_moshpit.filtering.tests': [
+            'data/*', 'data/mags/*', 'data/mags/*/*',
+        ],
     },
     zip_safe=False,
 )