annotate tests/test_plot.py @ 2:b6fd86c539ea

Fix pysam version in riboplot (was 0.7.7 instead of 0.8.3). Fix move sys import to the top in ribocount.py, riboplot.py. Minor: * Include an empty conditional when no RNA coverage is required. * Remove output directories after zip file is created.

import os
import logging
import unittest

from riboplot import core, config, plot

# use testing configuration
CONFIG = plot.CONFIG = 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 RNACountsTestCase(unittest.TestCase):

    def test_get_rna_counts(self):
        """Test get RNA counts for transcript from RNA-Seq BAM file"""
        counts = plot.get_rna_counts(rna_file=RNA_FILE, transcript_name=TRANSCRIPT_NAME)
        self.assertIsInstance(counts, dict)
        self.assertTrue(len(counts) > 0)

    def test_missing_rna_file(self):
        """Exit with error if RNA BAM file does not exist. """
        self.assertRaises(OSError, plot.get_rna_counts, rna_file='{}.absent'.format(RNA_FILE),
                          transcript_name=TRANSCRIPT_NAME)

    def test_missing_bedtools(self):
        """Exit with error if bedtools is missing."""
        # reset env temporarily
        paths = os.environ['PATH']
        os.environ['PATH'] = ''
        self.assertRaises(OSError, plot.get_rna_counts, rna_file=RNA_FILE,
                          transcript_name=TRANSCRIPT_NAME)
        os.environ['PATH'] = paths


class PlotTestCase(unittest.TestCase):

    def test_get_codon_positions(self):
        """Test get codon positions. """
        # input is the sequence obtained from get_transcript so no new lines.
        fasta = ('AACCGGAGCACCCAGAGAAAACCCACGCAAACGCAGGGAGAATTTGCAAACTCCACACA'
                 'GAAATGCCAGCTGATCCAGCCGAGCCTCGAGTCAGCATCCTTGCTTGTTGGATGCCTGA'
                 'TTGCAGTTCAACTCCAAACTCAGTTGGACCAGCTGATCAGTG')
        codon_positions = plot.get_start_stops(fasta)
        expected = {1: {'starts': [], 'stops': []},
                    2: {'starts': [], 'stops': [71, 116, 152]},
                    3: {'starts': [63, 111], 'stops': []}}
        self.assertEqual(codon_positions, expected)