from atropos.commands.trim.filters import (
    NContentFilter,
    DISCARD,
    KEEP,
    SingleWrapper,
    PairedWrapper,
)
from atropos.io.sequence import Sequence


def test_ncontentfilter():
    # third parameter is True if read should be discarded
    params = [
        ("AAA", 0, KEEP),
        ("AAA", 1, KEEP),
        ("AAACCTTGGN", 1, KEEP),
        ("AAACNNNCTTGGN", 0.5, KEEP),
        ("NNNNNN", 1, DISCARD),
        ("ANAAAA", 1 / 6, KEEP),
        ("ANAAAA", 0, DISCARD),
    ]
    for seq, count, expected in params:
        _filter = NContentFilter(count=count)
        _seq = Sequence("read1", seq, qualities="#" * len(seq))
        assert _filter(_seq) == expected


def test_ncontentfilter_paired():
    params = [
        ("AAA", "AAA", 0, KEEP),
        ("AAAN", "AAA", 0, DISCARD),
        ("AAA", "AANA", 0, DISCARD),
        ("ANAA", "AANA", 1, KEEP),
    ]
    for seq1, seq2, count, expected in params:
        _filter = NContentFilter(count=count)
        filter_legacy = SingleWrapper(_filter)
        filter_both = PairedWrapper(_filter)
        read1 = Sequence("read1", seq1, qualities="#" * len(seq1))
        read2 = Sequence("read1", seq2, qualities="#" * len(seq2))
        assert filter_legacy(read1, read2) == _filter(read1)
        # discard entire pair if one of the reads fulfills criteria
        assert filter_both(read1, read2) == expected
