0aceafbf |
1 | #!/usr/bin/python |
2 | |
3 | from __future__ import with_statement |
4 | from tempfile import mkdtemp |
5 | from shutil import rmtree |
6 | from itertools import imap |
7 | from string import hexdigits |
8 | import unittest, os |
9 | import pm3_eml2mfd, pm3_mfd2eml |
10 | |
11 | class TestEmlMfd(unittest.TestCase): |
12 | def setUp(self): |
13 | self.tmpdir = mkdtemp() |
14 | |
15 | def tearDown(self): |
16 | rmtree(self.tmpdir) |
17 | |
18 | EML2MFD_TESTCASES = [ |
19 | ('', ''), |
20 | ("41424344\r\n45464748\n494A4B4C\n", "ABCDEFGHIJKL") |
21 | ] |
22 | def test_eml2mfd(self): |
23 | self.three_argument_test(pm3_eml2mfd.main, self.EML2MFD_TESTCASES) |
24 | |
25 | def test_mfd2eml(self): |
26 | self.three_argument_test(pm3_mfd2eml.main, |
27 | imap(reversed, self.EML2MFD_TESTCASES), c14n=hex_c14n) |
28 | |
29 | def three_argument_test(self, operation, cases, c14n=str): |
30 | for case_input, case_output in cases: |
31 | try: |
32 | inp_name = os.path.join(self.tmpdir, 'input') |
33 | out_name = os.path.join(self.tmpdir, 'output') |
34 | with file(inp_name, 'wb') as in_file: |
35 | in_file.write(case_input) |
36 | operation(['', inp_name, out_name]) |
37 | with file(out_name, 'rb') as out_file: |
38 | self.assertEquals(c14n(case_output), c14n(out_file.read())) |
39 | finally: |
40 | for file_name in inp_name, out_name: |
41 | if os.path.exists(file_name): |
42 | os.remove(file_name) |
43 | |
44 | |
45 | def hex_c14n(inp): |
46 | """ |
47 | Canonicalizes the input string by removing non-hexadecimal |
48 | characters and making everything uppercase |
49 | """ |
50 | return ''.join(c.upper() for c in inp if c in hexdigits) |
51 | |
52 | if __name__ == '__main__': |
53 | unittest.main() |