# To run these tests:
#  - copy this into your assignment6 directory
#  - make sure you have a "codon_functions.py" file in that directory
#  - run this as "python test_codon_functions.py"
#
# The output should be "All tests passed."

from codon_functions import get_codons, translate_codons
from codon_functions import translate_codons_to_stop, translate_reading_frame

def test_get_codons():
    for seq, expected_codons in (
        # These lengths are a multiple of 3
        ("", []),
        ("ATC", ["ATC"]),
        ("ATCGAT", ["ATC", "GAT"]),
        ("ATCGTGCATAGACTATGCAATATACCG",
           ["ATC","GTG","CAT","AGA","CTA","TGC","AAT","ATA","CCG"]),

        # These lengths are not a multiple of 3
        ("A", []),
        ("TC", []),
        ("TTTC", ["TTT"]),
        ("TTTCC", ["TTT"]),
        ("AAATTTCCCGGGATCG", ["AAA", "TTT", "CCC", "GGG", "ATC"]),
        ):

        codons = list(get_codons(seq))
        if codons != expected_codons:
            raise AssertionError("Codons for %r was %r, expected %r" %
                                 (seq, codons, expected_codons))

    # This checks that the 'get_codon' function returns a generator
    gen = get_codons("ATC")
    if gen.next() != "ATC":
        raise AssertionError("Could not get the codon")
    try:
        gen.next()
    except StopIteration:
        pass
    else:
        raise AssertionError("Only supposed to find one codon")

def test_translate_codons():
    # A quick test to make sure the function allows lists of codons
    residues = list(translate_codons( ["GAG", "AAG", "TTG", "GCT", "GAT"]))
    if residues != ["E", "K", "L", "A", "D"]:
        raise AssertionError("Residues for list data was %r, expected %r" %
                             (residues, ["E", "K", "L", "A", "D"]))
    
    for (seq, expected_residues) in (
        # Multiple of 3, no stop codons, not ambiguous
        ("", []),
        ("GCT", ["A"]),
        ("GCTAAT", ["A", "N"]),
        ("TGGGAGCGTGATAATGCT", ["W", "E", "R", "D", "N", "A"]),
        
        # Not multiples of 3, no stop codons, not ambiguous
        ("A", []),
        ("GAGAAGTTGGCTGAT", ["E", "K", "L", "A", "D"]),
        
        # Test the stop codon
        ("TAATAGTGA", ["*", "*", "*"]),
        ("AAATAAAAATGAAAA", ["K", "*", "K", "*", "K"]),
        
        # Test the ambiguous codes
        ("ATH", ["I"]),
        ("AGRAGY", ["R", "S"]),
        ("AAAAGRTAGATHC", ["K", "R", "*", "I"]),
        ):
        
        residues = list(translate_codons(get_codons(seq)))
        if residues != expected_residues:
            raise AssertionError("Residues for %r was %r, expected %r" %
                                 (seq, residues, expected_residues))
    
    # This checks that the 'translate_codons' function uses an iterator
    # and returns a generator
    def yield_codons():
        yield "GCT"
        raise AssertionError("I didn't ask for the second codon")

    gen = translate_codons(yield_codons())
    if gen.next() != "A":
        raise AssertionError("Could not get the codon")
    try:
        # Make sure it raises the expected
        gen.next()
    except AssertionError:
        pass
    else:
        raise AssertionError("There was a second item?")

def test_translate_codons_to_stop():
    for (seq, expected_residues) in (
        # Multiple of 3, no stop codons, not ambiguous
        ("", []),
        ("GCT", ["A"]),
        ("GCTAAT", ["A", "N"]),
        ("TGGGAGCGTGATAATGCT", ["W", "E", "R", "D", "N", "A"]),
        
        # Not multiples of 3, no stop codons, not ambiguous
        ("A", []),
        ("GAGAAGTTGGCTGAT", ["E", "K", "L", "A", "D"]),
        
        # Test the stop codon
        ("TAATAGTGA", []),
        ("AAATAAAAATGAAAA", ["K"]),
        ("GAGAAATGA", ["E", "K"]),
        
        # Test the ambiguous codes
        ("ATH", ["I"]),
        ("AGRAGY", ["R", "S"]),
        ("AAAAGRTAGATHC", ["K", "R"]),
        ):
        
        residues = list(translate_codons_to_stop(get_codons(seq)))
        if residues != expected_residues:
            raise AssertionError("Residues for %r was %r, expected %r" %
                                 (seq, residues, expected_residues))
    
    # This checks that the 'translate_codons_to_stop' function uses an iterator
    # and returns a generator
    def yield_codons():
        yield "GCT"
        raise AssertionError("I didn't ask for the second codon")

    gen = translate_codons_to_stop(yield_codons())
    if gen.next() != "A":
        raise AssertionError("Could not get the codon")
    try:
        # Make sure it raises the expected
        gen.next()
    except AssertionError:
        pass
    else:
        raise AssertionError("There was a second item?")

def test_translate_reading_frame():
    for (seq, expected_orf) in (
        ("", []),
        
        # Has a start and an end
        ("CTGAAATGA", ["L", "K"]),
        # start codon is the 3rd residue
        ("ATHAGYCTGAAATGA", ["L", "K"]),
        # make sure it does not read multiple reading frames
        ("ATHAGYCTGAAATGAATHAGYCTGAAATGA", ["L", "K"]),

        # Test a different start codon; go to the end (no stop codon)
        ("AAAATGGGGGCT", ["M", "G", "A"]),

        # Make sure only one start codon is accepted
        ("AAAATGGGGATGGCT", ["M", "G", "M", "A"]),
        ):

        orf = list(translate_reading_frame(get_codons(seq)))
        if orf != expected_orf:
            raise AssertionError("ORF for %r was %r, expected %r" %
                                 (seq, orf, expected_orf))

    # Check that the translate_reading_frame function works with
    # both lists and iterables
    orf = translate_reading_frame(["ATH", "AGY", "CTG", "AAA", "TGA"])
    x = orf.next()
    if x != "L":
        raise AssertionError("First residue is %r" % (x,))
    x = orf.next()
    if x != "K":
        raise AssertionError("Second residue is %r" % (x,))
    try:
        x = orf.next()
    except StopIteration:
        pass
    else:
        raise AssertionError("Unexpected third residue %r" % (x,))

def test():
    test_get_codons()
    test_translate_codons()
    test_translate_codons_to_stop()
    test_translate_reading_frame()
        
if __name__ == "__main__":
    test()
    print "All tests passed."