Added command line arguments tests.

Utils/Dataflow/pyDKB/dataflow/stage/tests/test_ProcessorStage.py

@@ -0,0 +1,137 @@
+#!/usr/bin/env python
+
+"""
+Tests for pyDKB.dataflow.stage.ProcessorStage.
+Usage: 'python -m unittest discover' from ..
+(directory with pyDKB.dataflow.stage code).
+"""
+
+
+import os
+import sys
+import unittest
+
+
+# Relative import inside of pyDKB prevents the use of simple 'import pyDKB'.
+try:
+    base_dir = os.path.dirname(__file__)  # Directory with this file
+    dkb_dir = os.path.join(base_dir, os.pardir)  # stage directory
+    dkb_dir = os.path.join(dkb_dir, os.pardir)  # dataflow directory
+    dkb_dir = os.path.join(dkb_dir, os.pardir)  # pyDKB's directory
+    dkb_dir = os.path.join(dkb_dir, os.pardir)  # pyDKB's parent directory
+    sys.path.append(dkb_dir)
+    import pyDKB
+except Exception, err:
+    sys.stderr.write("(ERROR) Failed to import pyDKB library: %s\n" % err)
+    sys.exit(1)
+
+
+class ProcessorStageArgsTestCase(unittest.TestCase):
+    default_args = {
+        'mode': 'f',
+        'config': None,
+        'eom': '\n',
+        'eop': '',
+        'source': 'f',
+        'dest': 'f',
+        'input_dir': os.curdir,
+        'output_dir': 'out',
+        'hdfs': False,
+        'input_files': []
+    }
+
+    def setUp(self):
+        self.stage = pyDKB.dataflow.stage.ProcessorStage()
+
+    def tearDown(self):
+        self.stage = None
+
+    def check_args(self, args):
+        for a in args:
+            self.assertEqual(getattr(self.stage.ARGS, a), args[a])
+
+    def test_default(self):
+        self.stage.parse_args('')
+        self.check_args(self.default_args)
+
+    def test_hdfs(self):
+        self.stage.parse_args(['--hdfs'])
+        args = dict(self.default_args)
+        args['hdfs'] = True
+        self.check_args(args)
+
+    def test_eom(self):
+        self.stage.parse_args(['-e', '\t'])
+        args = dict(self.default_args)
+        args['eom'] = '\t'
+        self.check_args(args)
+
+    def test_eop(self):
+        self.stage.parse_args(['-E', '\t'])
+        args = dict(self.default_args)
+        args['eop'] = '\t'
+        self.check_args(args)
+
+    def test_input_dir(self):
+        self.stage.parse_args(['--input-dir', 'something'])
+        args = dict(self.default_args)
+        args['input_dir'] = 'something'
+        self.check_args(args)
+
+    def test_i(self):
+        self.stage.parse_args(['-i', 'something'])
+        args = dict(self.default_args)
+        args['input_dir'] = 'something'
+        self.check_args(args)
+
+    def test_output_dir(self):
+        self.stage.parse_args(['--output-dir', 'something'])
+        args = dict(self.default_args)
+        args['output_dir'] = 'something'
+        self.check_args(args)
+
+    def test_o(self):
+        self.stage.parse_args(['-o', 'something'])
+        args = dict(self.default_args)
+        args['output_dir'] = 'something'
+        self.check_args(args)
+
+
+args_to_add = {
+    'source': ['f', 's', 'h'],
+    'dest': ['f', 's', 'h'],
+    'mode': ['f', 's', 'm'],
+}
+
+
+def add_arg(arg, val, short=False):
+    def f(self):
+        if short:
+            self.stage.parse_args(['-' + arg[0], val])
+        else:
+            self.stage.parse_args(['--' + arg, val])
+        args = dict(self.default_args)
+        args[arg] = val
+        self.check_args(args)
+    if short:
+        setattr(ProcessorStageArgsTestCase, 'test_%s_%s' % (arg[0], val), f)
+    else:
+        setattr(ProcessorStageArgsTestCase, 'test_%s_%s' % (arg, val), f)
+
+
+for a in args_to_add:
+    for v in args_to_add[a]:
+        add_arg(a, v)
+        add_arg(a, v, True)
+
+
+test_cases = (
+    ProcessorStageArgsTestCase,
+)
+
+
+def load_tests(loader, tests, pattern):
+    suite = unittest.TestSuite()
+    for case in test_cases:
+        suite.addTest(loader.loadTestsFromTestCase(case))
+    return suite