comparison tests/test_riboplot.py @ 3:8e1efafa6277

Updated version * Bugfix: blue lines in some plots (bar colors were not set correctly) * Cookiecutter template * Additional unit tests * Add plot legend

3:8e1efafa6277

import os
import shutil
import logging
import unittest
import tempfile

from riboplot import ribocore, riboplot

# use testing configuration
CONFIG = riboplot.CONFIG = riboplot.config.TestingConfig()
logging.disable(logging.CRITICAL)

RIBO_FILE = os.path.join(CONFIG.DATA_DIR, '5hRPFsorted.bam')
RNA_FILE = os.path.join(CONFIG.DATA_DIR, '5hmRNAsorted.bam')
TRANSCRIPT_NAME = 'gi|148357119|ref|NM_001098396.1|'
TRANSCRIPTOME_FASTA = os.path.join(CONFIG.DATA_DIR, 'zebrafish.fna')
TRANSCRIPTOME_FASTA_MINUS1 = os.path.join(CONFIG.DATA_DIR, 'zebrafish_minus1.fna')


class CheckArgumentsTestCase(unittest.TestCase):
    """Check if all arguments sent on the command line are valid."""
    parser = riboplot.create_parser()

    def test_bedtools_missing(self):
        """If bedtools is not in PATH, raise an error."""
        args = self.parser.parse_args(
            ['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', TRANSCRIPT_NAME, '-n', RNA_FILE])
        save_path = os.environ['PATH']
        os.environ['PATH'] = ''
        self.assertRaises(OSError, ribocore.check_optional_arguments, args.ribo_file, args.rna_file)
        os.environ['PATH'] = save_path

    def test_is_bam_valid(self):
        """Test if BAM file is valid."""
        valid = ribocore.is_bam_valid(RIBO_FILE)
        self.assertTrue(valid)

        # test with a FASTA file (which is not BAM)
        self.assertRaises(ValueError, ribocore.is_bam_valid, TRANSCRIPTOME_FASTA)

    def test_bam_has_index(self):
        """Check if BAM file has an index."""
        # RPF file has an index
        has_index = ribocore.bam_has_index(RIBO_FILE)
        self.assertTrue(has_index)

        # RNA file doesn't have an index
        has_index = ribocore.bam_has_index(RNA_FILE)
        self.assertFalse(has_index)

    def test_create_bam_index(self):
        """Index a BAM file."""
        ribocore.create_bam_index(RNA_FILE)

        # check if index exists
        has_index = ribocore.bam_has_index(RNA_FILE)
        self.assertTrue(has_index)

        # remove index
        os.remove('{}.bai'.format(RNA_FILE))

    def test_valid_read_length(self):
        """Read length should be a valid integer."""
        args = self.parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA,
                                       '-t', TRANSCRIPT_NAME, '-l', '28'])
        ribocore.check_optional_arguments(ribo_file=args.ribo_file, read_length=args.read_length)

    def test_invalid_read_length(self):
        """An error is raised if an invalid read length is used."""
        args = self.parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', TRANSCRIPT_NAME,
                                       '-l', '-1'])  # invalid read length -1
        self.assertRaises(ribocore.ArgumentError, ribocore.check_optional_arguments,
                          args.ribo_file, None, args.read_length)

        args = self.parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', TRANSCRIPT_NAME,
                                       '-l', '100'])  # invalid read length 100
        self.assertRaises(ribocore.ArgumentError, ribocore.check_optional_arguments,
                          args.ribo_file, None, args.read_length)

    def test_valid_read_offset(self):
        """Read offset should be positive."""
        args = self.parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', TRANSCRIPT_NAME,
                                       '-s', '-1'])  # invalid read offset -1
        self.assertRaises(ribocore.ArgumentError, ribocore.check_optional_arguments,
                          args.ribo_file, None, None, args.read_offset)

    def test_is_fasta_valid(self):
        """A valid FASTA file can be opened with pysam.FastaFile."""
        self.assertTrue(ribocore.is_fasta_valid(TRANSCRIPTOME_FASTA))

    def test_missing_transcript_in_fasta(self):
        """If a transcript is missing in FASTA, an error is raised."""
        args = self.parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', TRANSCRIPT_NAME])  # invalid read offset -1
        self.assertRaises(ribocore.ArgumentError, ribocore.check_required_arguments,
                          args.ribo_file, args.transcriptome_fasta, 'hello')

    def test_missing_transcript_in_bam(self):
        """If a transcript is missing in BAM, an error is raised."""
        # testing with an unrelated BAM file
        args = self.parser.parse_args(['-b', '/home/vimal/tmp/empty_tp/RiboSeq.bam', '-f', TRANSCRIPTOME_FASTA,
                                       '-t', TRANSCRIPT_NAME])
        self.assertRaises(ribocore.ArgumentError, ribocore.check_required_arguments, args.ribo_file,
                          args.transcriptome_fasta, args.transcript_name)


class RNACountsTestCase(unittest.TestCase):

    def test_get_rna_counts(self):
        """Test get RNA counts for transcript from RNA-Seq BAM file. Assumes bedtools is installed."""
        counts = riboplot.get_rna_counts(RNA_FILE, TRANSCRIPT_NAME)
        self.assertIsInstance(counts, dict)
        self.assertTrue(len(counts) > 0)


class RiboPlotTestCase(unittest.TestCase):

    def test_get_codon_positions(self):
        """Get codon positions in all frames given a sequence."""
        # the positions on this sequence were calculated manually.
        fasta = ('AACCGGAGCACCCAGAGAAAACCCACGCAAACGCAGGGAGAATTTGCAAACTCCACACA'
                 'GAAATGCCAGCTGATCCAGCCGAGCCTCGAGTCAGCATCCTTGCTTGTTGGATGCCTGA'
                 'TTGCAGTTCAACTCCAAACTCAGTTGGACCAGCTGATCAGTG')
        codon_positions = riboplot.get_start_stops(fasta)
        expected = {1: {'starts': [], 'stops': []},
                    2: {'starts': [], 'stops': [71, 116, 152]},
                    3: {'starts': [63, 111], 'stops': []}}
        self.assertEqual(codon_positions, expected)

    def test_valid_riboplot_run(self):
        """A good riboplot run"""
        output_dir = tempfile.mkdtemp()
        print 'Output path is {}'.format(output_dir)
        parser = riboplot.create_parser()
        args = parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', TRANSCRIPT_NAME,
                                  '-o', output_dir])
        riboplot.main(args)
        for fname in ('riboplot.png', 'riboplot.svg', 'RiboCounts.csv'):
            self.assertTrue(os.path.exists(os.path.join(output_dir, fname)))
        shutil.rmtree(output_dir)

    def test_transcript_with_no_counts(self):
        """If the transcript has no ribocounts, no plot should be produced."""
        transcript = 'gi|62955616|ref|NM_001017822.1|' # has no reads
        output_dir = tempfile.mkdtemp()
        parser = riboplot.create_parser()
        args = parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', transcript, '-o', output_dir])
        self.assertRaises(ribocore.RiboPlotError, riboplot.main, args)
        for fname in ('riboplot.png', 'riboplot.svg', 'RiboCounts.csv'):
            self.assertFalse(os.path.exists(os.path.join(output_dir, fname)))
        shutil.rmtree(output_dir)

    @unittest.skip('todo')
    def test_get_ribo_counts(self):
        """Get RiboSeq read counts"""
        pass

    @unittest.skip('todo')
    def test_write_ribo_counts(self):
        """Write RiboSeq read counts as CSV."""
        pass

    @unittest.skip('todo')
    def test_plot_read_counts(self):
        """Generate riboplots"""
        pass