Mercurial > repos > vimalkumarvelayudhan > riboplot

comparison tests/test_riboplot.py @ 3:8e1efafa6277

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Updated version * Bugfix: blue lines in some plots (bar colors were not set correctly) * Cookiecutter template * Additional unit tests * Add plot legend
author Vimalkumar Velayudhan <vimal@biotechcoder.com>
date Wed, 12 Aug 2015 09:27:45 +0100
parents
children 2ffa8172dce1
comparison
equal deleted inserted replaced
2:b6fd86c539ea 3:8e1efafa6277
1 import os
2 import shutil
3 import logging
4 import unittest
5 import tempfile
6
7 from riboplot import ribocore, riboplot
8
9 # use testing configuration
10 CONFIG = riboplot.CONFIG = riboplot.config.TestingConfig()
11 logging.disable(logging.CRITICAL)
12
13 RIBO_FILE = os.path.join(CONFIG.DATA_DIR, '5hRPFsorted.bam')
14 RNA_FILE = os.path.join(CONFIG.DATA_DIR, '5hmRNAsorted.bam')
15 TRANSCRIPT_NAME = 'gi|148357119|ref|NM_001098396.1|'
16 TRANSCRIPTOME_FASTA = os.path.join(CONFIG.DATA_DIR, 'zebrafish.fna')
17 TRANSCRIPTOME_FASTA_MINUS1 = os.path.join(CONFIG.DATA_DIR, 'zebrafish_minus1.fna')
18
19
20 class CheckArgumentsTestCase(unittest.TestCase):
21 """Check if all arguments sent on the command line are valid."""
22 parser = riboplot.create_parser()
23
24 def test_bedtools_missing(self):
25 """If bedtools is not in PATH, raise an error."""
26 args = self.parser.parse_args(
27 ['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', TRANSCRIPT_NAME, '-n', RNA_FILE])
28 save_path = os.environ['PATH']
29 os.environ['PATH'] = ''
30 self.assertRaises(OSError, ribocore.check_optional_arguments, args.ribo_file, args.rna_file)
31 os.environ['PATH'] = save_path
32
33 def test_is_bam_valid(self):
34 """Test if BAM file is valid."""
35 valid = ribocore.is_bam_valid(RIBO_FILE)
36 self.assertTrue(valid)
37
38 # test with a FASTA file (which is not BAM)
39 self.assertRaises(ValueError, ribocore.is_bam_valid, TRANSCRIPTOME_FASTA)
40
41 def test_bam_has_index(self):
42 """Check if BAM file has an index."""
43 # RPF file has an index
44 has_index = ribocore.bam_has_index(RIBO_FILE)
45 self.assertTrue(has_index)
46
47 # RNA file doesn't have an index
48 has_index = ribocore.bam_has_index(RNA_FILE)
49 self.assertFalse(has_index)
50
51 def test_create_bam_index(self):
52 """Index a BAM file."""
53 ribocore.create_bam_index(RNA_FILE)
54
55 # check if index exists
56 has_index = ribocore.bam_has_index(RNA_FILE)
57 self.assertTrue(has_index)
58
59 # remove index
60 os.remove('{}.bai'.format(RNA_FILE))
61
62 def test_valid_read_length(self):
63 """Read length should be a valid integer."""
64 args = self.parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA,
65 '-t', TRANSCRIPT_NAME, '-l', '28'])
66 ribocore.check_optional_arguments(ribo_file=args.ribo_file, read_length=args.read_length)
67
68 def test_invalid_read_length(self):
69 """An error is raised if an invalid read length is used."""
70 args = self.parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', TRANSCRIPT_NAME,
71 '-l', '-1']) # invalid read length -1
72 self.assertRaises(ribocore.ArgumentError, ribocore.check_optional_arguments,
73 args.ribo_file, None, args.read_length)
74
75 args = self.parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', TRANSCRIPT_NAME,
76 '-l', '100']) # invalid read length 100
77 self.assertRaises(ribocore.ArgumentError, ribocore.check_optional_arguments,
78 args.ribo_file, None, args.read_length)
79
80 def test_valid_read_offset(self):
81 """Read offset should be positive."""
82 args = self.parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', TRANSCRIPT_NAME,
83 '-s', '-1']) # invalid read offset -1
84 self.assertRaises(ribocore.ArgumentError, ribocore.check_optional_arguments,
85 args.ribo_file, None, None, args.read_offset)
86
87 def test_is_fasta_valid(self):
88 """A valid FASTA file can be opened with pysam.FastaFile."""
89 self.assertTrue(ribocore.is_fasta_valid(TRANSCRIPTOME_FASTA))
90
91 def test_missing_transcript_in_fasta(self):
92 """If a transcript is missing in FASTA, an error is raised."""
93 args = self.parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', TRANSCRIPT_NAME]) # invalid read offset -1
94 self.assertRaises(ribocore.ArgumentError, ribocore.check_required_arguments,
95 args.ribo_file, args.transcriptome_fasta, 'hello')
96
97 def test_missing_transcript_in_bam(self):
98 """If a transcript is missing in BAM, an error is raised."""
99 # testing with an unrelated BAM file
100 args = self.parser.parse_args(['-b', '/home/vimal/tmp/empty_tp/RiboSeq.bam', '-f', TRANSCRIPTOME_FASTA,
101 '-t', TRANSCRIPT_NAME])
102 self.assertRaises(ribocore.ArgumentError, ribocore.check_required_arguments, args.ribo_file,
103 args.transcriptome_fasta, args.transcript_name)
104
105
106 class RNACountsTestCase(unittest.TestCase):
107
108 def test_get_rna_counts(self):
109 """Test get RNA counts for transcript from RNA-Seq BAM file. Assumes bedtools is installed."""
110 counts = riboplot.get_rna_counts(RNA_FILE, TRANSCRIPT_NAME)
111 self.assertIsInstance(counts, dict)
112 self.assertTrue(len(counts) > 0)
113
114
115 class RiboPlotTestCase(unittest.TestCase):
116
117 def test_get_codon_positions(self):
118 """Get codon positions in all frames given a sequence."""
119 # the positions on this sequence were calculated manually.
120 fasta = ('AACCGGAGCACCCAGAGAAAACCCACGCAAACGCAGGGAGAATTTGCAAACTCCACACA'
121 'GAAATGCCAGCTGATCCAGCCGAGCCTCGAGTCAGCATCCTTGCTTGTTGGATGCCTGA'
122 'TTGCAGTTCAACTCCAAACTCAGTTGGACCAGCTGATCAGTG')
123 codon_positions = riboplot.get_start_stops(fasta)
124 expected = {1: {'starts': [], 'stops': []},
125 2: {'starts': [], 'stops': [71, 116, 152]},
126 3: {'starts': [63, 111], 'stops': []}}
127 self.assertEqual(codon_positions, expected)
128
129 def test_valid_riboplot_run(self):
130 """A good riboplot run"""
131 output_dir = tempfile.mkdtemp()
132 print 'Output path is {}'.format(output_dir)
133 parser = riboplot.create_parser()
134 args = parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', TRANSCRIPT_NAME,
135 '-o', output_dir])
136 riboplot.main(args)
137 for fname in ('riboplot.png', 'riboplot.svg', 'RiboCounts.csv'):
138 self.assertTrue(os.path.exists(os.path.join(output_dir, fname)))
139 shutil.rmtree(output_dir)
140
141 def test_transcript_with_no_counts(self):
142 """If the transcript has no ribocounts, no plot should be produced."""
143 transcript = 'gi|62955616|ref|NM_001017822.1|' # has no reads
144 output_dir = tempfile.mkdtemp()
145 parser = riboplot.create_parser()
146 args = parser.parse_args(['-b', RIBO_FILE, '-f', TRANSCRIPTOME_FASTA, '-t', transcript, '-o', output_dir])
147 self.assertRaises(ribocore.RiboPlotError, riboplot.main, args)
148 for fname in ('riboplot.png', 'riboplot.svg', 'RiboCounts.csv'):
149 self.assertFalse(os.path.exists(os.path.join(output_dir, fname)))
150 shutil.rmtree(output_dir)
151
152 @unittest.skip('todo')
153 def test_get_ribo_counts(self):
154 """Get RiboSeq read counts"""
155 pass
156
157 @unittest.skip('todo')
158 def test_write_ribo_counts(self):
159 """Write RiboSeq read counts as CSV."""
160 pass
161
162 @unittest.skip('todo')
163 def test_plot_read_counts(self):
164 """Generate riboplots"""
165 pass