add unit test and update argparsing

1 files changed, 45 insertions, 25 deletions

diff --git a/calctenna.py b/calctenna.py
index b782e5c..34cc2b8 100755
--- a/calctenna.py
+++ b/calctenna.py
@@ -6,23 +6,21 @@
 # Oliver Meckmann
 
 import argparse
+import unittest
+import sys
+
+class Tests(unittest.TestCase):
+    def test_calc(self):
+        self.assertEqual(calc(100e6, 0.951),
+                (2.99792458, 1.49896229, 0.749481145, 0.7127565688949999))
 
 def run_argparse():
     argparser = argparse.ArgumentParser()
-    argparser.add_argument(
-        '-f',
-        '--frequency',
-        required=True,
-        metavar='frequency',
-        help='desired frequency in Hz (e.g. 144800000 or 433e6, ...)',
-        nargs='?',
-        #default='100e6'
-    )
+
     argparser.add_argument(
         '-V',
         '--velocity-factor',
         metavar='velocity_factor',
-        nargs='?',
         default='0.951'
     )
     argparser.add_argument(
@@ -31,47 +29,69 @@ def run_argparse():
         help='print only result',
         action='store_true'
     )
+    argparser.add_argument(
+        'frequency',
+        metavar='frequency',
+        help='desired frequency in Hz (e.g. 144800000 or 433e6, ...)',
+        nargs='?'
+    )
+    argparser.add_argument(
+        '-t',
+        '--tests',
+        help='run unit tests and exit',
+        action='store_true'
+    )
+
     args = argparser.parse_args()
 
     if not args.frequency:
         argparser.print_usage()
-        exit()
 
     return args
 
-def calc(args):
-    c = 299792458 # speed of light
-
+def check_input(frequency, velocity_factor):
     try:
-        f = float(args.frequency)
+        f = float(frequency)
     except:
-        print("Sorry, frequency input invalid")
-        exit()
+        print("Sorry, frequency input invalid or missing")
+        return False
 
     if f <= 0.0:
         print("Sorry, frequency must be > 0")
-        exit()
+        return False
 
     try:
-        V = float(args.velocity_factor)
+        V = float(velocity_factor)
     except:
         print("Sorry, velocity factor input invalid")
-        exit()
+        return False
 
     if V < 0.0 or V > 1.0:
         print("Sorry, velocity factor can only be between 0.0 and 1.0")
-        exit()
+        return False
 
-    wavelength = float(c) / f
+    return True
+
+def calc(frequency, velocity_factor):
+    c = 299792458 # speed of light
+    wavelength = float(c) / float(frequency)
     wavelength_2 = wavelength/2
     wavelength_4 = wavelength/4
-    dipole_leg_length = wavelength/4 * V
-
+    dipole_leg_length = wavelength/4 * float(velocity_factor)
     return (wavelength, wavelength_2, wavelength_4, dipole_leg_length)
 
 if __name__ == '__main__':
     args = run_argparse()
-    (wavelength, wavelength_2, wavelength_4, dipole_leg_length) = calc(args)
+
+    if args.tests:
+        # override args for unittest: argv[0] to have name in test output
+        unittest.main(argv=[sys.argv[0], '--verbose'], exit=True)
+
+    if not check_input(args.frequency, args.velocity_factor):
+        exit()
+
+    (wavelength, wavelength_2, wavelength_4, dipole_leg_length) = \
+            calc(args.frequency, args.velocity_factor)
 
     if not args.quiet:
         print("Desired Frequency: " + str(args.frequency) + " Hz")