--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/m209/tests/test_main.py	Wed Jul 03 20:55:20 2013 -0500
@@ -0,0 +1,136 @@
+# Copyright (C) 2013 by Brian Neal.
+# This file is part of m209, the M-209 simulation.
+# m209 is released under the MIT License (see LICENSE.txt).
+
+import os
+import tempfile
+import unittest
+
+from ..main import main
+
+
+class KeyGenTestCase(unittest.TestCase):
+
+    def setUp(self):
+
+        self.fp = tempfile.NamedTemporaryFile(mode='w')
+
+    def tearDown(self):
+
+        self.fp.close()
+
+    def test_overwrite(self):
+        """Verify we exit if the key file already exists"""
+
+        argv = ['keygen', '--start=YA', '--number=100', '-z', self.fp.name]
+        self.assertRaises(SystemExit, main, argv)
+
+    def test_too_many(self):
+        """Verify we exit if we can't generate N key lists if start is too
+        high
+
+        """
+        argv = ['keygen', '--start=YA', '--number=100', '-o', '-z', self.fp.name]
+        self.assertRaises(SystemExit, main, argv)
+
+    def test_nominal_random(self):
+        """Test we can generate N key lists with random indicators"""
+
+        argv = ['keygen', '--number=10', '-o', '-z', self.fp.name]
+        main(argv)
+
+    def test_nominal_start(self):
+        """Test we can generate N key lists with a fixed starting indicator"""
+
+        argv = ['keygen', '--start=GG', '--number=10', '-o', '-z', self.fp.name]
+        main(argv)
+
+
+class EncryptDecryptBadArgsTestCase(unittest.TestCase):
+
+    def test_no_key_file(self):
+        """Ensure we exit if key file doesn't exist"""
+
+        fp = tempfile.NamedTemporaryFile()
+        name = fp.name
+        fp.close()
+
+        argv = ['encrypt', '--text=TEST', '-z', name]
+        self.assertRaises(SystemExit, main, argv)
+
+        argv = ['decrypt', '--text=TEST', '-z', name]
+        self.assertRaises(SystemExit, main, argv)
+
+
+class EncryptDecryptTestCase(unittest.TestCase):
+
+    def setUp(self):
+
+        self.fp = tempfile.NamedTemporaryFile(mode='w')
+        argv = ['keygen', '--start=GG', '--number=10', '-o', '-z', self.fp.name]
+        main(argv)
+
+    def tearDown(self):
+
+        self.fp.close()
+
+    def test_conflicting_sources(self):
+        """Ensure -f or -t is supplied but not both"""
+
+        argv = ['encrypt', '--text=TEST', '-f', '-', '-z', self.fp.name]
+        self.assertRaises(SystemExit, main, argv)
+
+        argv = ['encrypt', '-z', self.fp.name]
+        self.assertRaises(SystemExit, main, argv)
+
+        argv = ['decrypt', '--text=TEST', '-f', '-', '-z', self.fp.name]
+        self.assertRaises(SystemExit, main, argv)
+
+        argv = ['decrypt', '-z', self.fp.name]
+        self.assertRaises(SystemExit, main, argv)
+
+    def test_encrypt_text(self):
+
+        argv = ['encrypt', '--text=TEST', '--key-list-ind=GG', '-z', self.fp.name]
+        main(argv)
+
+    def test_encrypt_text_no_key_list(self):
+
+        argv = ['encrypt', '--text=TEST', '--key-list-ind=GA', '-z', self.fp.name]
+        self.assertRaises(SystemExit, main, argv)
+
+    def test_encrypt_file(self):
+
+        infile = tempfile.NamedTemporaryFile(mode='w', delete=False)
+        infile.write("TEST")
+        filename = infile.name
+        infile.close()
+
+        argv = ['encrypt', '-f', filename, '--key-list-ind=GG', '-z', self.fp.name]
+        try:
+            main(argv)
+        finally:
+            os.remove(filename)
+
+    def test_decrypt_text(self):
+
+        argv = ['decrypt', '-t', 'OOOZS IENGA DSGJX OOOZS IENGA', '-z', self.fp.name]
+        self.assertRaises(SystemExit, main, argv)
+
+    def test_decrypt_text_no_key_list(self):
+
+        argv = ['decrypt', '-t', 'OOOZS IENGG DSGJX OOOZS IENGG', '-z', self.fp.name]
+        main(argv)
+
+    def test_decrypt_file(self):
+
+        infile = tempfile.NamedTemporaryFile(mode='w', delete=False)
+        infile.write("OOOZS IENGG DSGJX OOOZS IENGG")
+        filename = infile.name
+        infile.close()
+
+        argv = ['decrypt', '-f', filename, '-z', self.fp.name]
+        try:
+            main(argv)
+        finally:
+            os.remove(filename)