separate test configuration logic from tests

bin/test_stolos.sh

@@ -8,7 +8,7 @@ echo $JOB_ID_VALIDATIONS
 echo $CONFIGURATION_BACKEND
 
 echo -n Is a local Zookeeper Server running and available?
-ans=`pgrep -f 'java -Dzookeeper' >/dev/null && echo yes || echo no`
+ans=`pgrep -f '-Dzookeeper' >/dev/null && echo yes || echo no`
 echo ...$ans
 echo -n Is a local Redis Server running and available?
 ans2=`pgrep -f 'redis-server' && echo yes || echo no`

stolos/testing_tools.py

@@ -1,6 +1,20 @@
 """
 Useful utilities for testing
 """
+from contextlib import contextmanager
+import os
+from os.path import join
+import tempfile
+import ujson
+import nose.tools as nt
+
+from stolos import dag_tools as dt
+from stolos.configuration_backend.json_config import JSONMapping
+from stolos import zookeeper_tools as zkt
+
+
+TASKS_JSON = os.environ['TASKS_JSON']
+TASKS_JSON_TMPFILES = {}
 
 
 def configure_logging(logname):
@@ -31,6 +45,7 @@ def format(self, record):
         _h.setFormatter(_formatter)
         log.addHandler(_h)
     log.setLevel(logging.DEBUG)
+    log.propagate = False
     return logging.getLogger(logname)
 
 
@@ -45,18 +60,18 @@ def _smart_run(func, args, kwargs):
     func(*args2, **kwargs2)
 
 
-def with_setup(setup=None, teardown=None, params=False):
-    """Decorator to add setup and/or teardown methods to a test function::
+def _with_setup(setup=None, teardown=None, params=False):
+    """Decorator to add setup and/or teardown methods to a test function
 
       @with_setup(setup, teardown)
       def test_something():
           " ... "
 
-    Note that `with_setup` is useful *only* for test functions, not for test
-    methods or inside of TestCase subclasses.
+    This extends the nose with_setup decorator such that:
+        - `setup` can return args and kwargs.
+        - The decorated test function and `teardown` can use any of the args
+          and kwargs returned by setup
     """
-    import nose.tools as nt
-
     def decorate(func, setup=setup, teardown=teardown):
         args = []
         kwargs = {}
@@ -68,7 +83,10 @@ def func_wrapped():
             func_wrapped = func
         if setup:
             def setup_wrapped():
-                rv = setup()
+                if params:
+                    rv = setup(func.func_name)
+                else:
+                    rv = setup()
                 if rv is not None:
                     args.extend(rv[0])
                     kwargs.update(rv[1])
@@ -96,9 +114,215 @@ def _t():
                 func_wrapped.teardown = _t
             else:
                 if params:
-                    func_wrapped.teardown = lambda: _smart_run(
-                        teardown, args, kwargs)
+                    def _sr():
+                        return _smart_run(teardown, args, kwargs)
+                    func_wrapped.teardown = _sr
                 else:
                     func_wrapped.teardown = teardown
         return func_wrapped
     return decorate
+
+
+def _create_tasks_json(fname_suffix='', inject={}, rename=False):
+    """
+    Create a new default tasks.json file
+    This of this like a useful setup() operation before tests.
+
+    `inject` - (dct) (re)define new tasks to add to the tasks graph
+    mark this copy as the new default.
+    `fname_suffix` - suffix in the file name of the new tasks.json file
+    `rename` - if True, change the name all tasks to include the fname_suffix
+
+    """
+    tasks_config = dt.get_tasks_config().cache
+    tasks_config.update(inject)  # assume we're using a json config
+
+    f = tempfile.mkstemp(prefix='tasks_json', suffix=fname_suffix)[1]
+    frv = ujson.dumps(tasks_config)
+    if rename:
+        renames = [(k, "%s__%s" % (k, fname_suffix))
+                   for k in tasks_config]
+        for k, new_k in renames:
+            frv = frv.replace(ujson.dumps(k), ujson.dumps(new_k))
+    with open(f, 'w') as fout:
+        fout.write(frv)
+
+    os.environ['TASKS_JSON'] = f
+    try:
+        dt.build_dag_cached.cache.clear()
+    except AttributeError:
+        pass
+    return f
+
+
+def _setup_func(func_name):
+    """Code that runs just before each test and configures a tasks.json file
+    for each test.  The tasks.json tmp files are stored in a global var."""
+    zk = zkt.get_zkclient('localhost:2181')
+    zk.delete('test_stolos', recursive=True)
+
+    rv = dict(
+        zk=zk,
+        oapp1='test_stolos/test_app',
+        oapp2='test_stolos/test_app2',
+        app3='test_stolos/test_app3__%s' % func_name,
+        app4='test_stolos/test_app4__%s' % func_name,
+        job_id1='20140606_1111_profile',
+        job_id2='20140606_2222_profile',
+        job_id3='20140604_1111_profile',
+        depends_on1='test_stolos/test_depends_on__%s' % func_name,
+        bash1='test_stolos/test_bash__%s' % func_name,
+        func_name=func_name,
+    )
+    rv.update(
+        app1='%s__%s' % (rv['oapp1'], func_name),
+        app2='%s__%s' % (rv['oapp2'], func_name),
+    )
+
+    f = _create_tasks_json(fname_suffix=func_name, rename=True)
+    TASKS_JSON_TMPFILES[func_name] = f
+    return ((), rv)
+
+
+def _teardown_func(func_name):
+    """Code that runs just after each test"""
+    os.environ['TASKS_JSON'] = TASKS_JSON
+    os.remove(TASKS_JSON_TMPFILES[func_name])
+
+
+def with_setup(func, setup_func=_setup_func, teardown_func=_teardown_func):
+    """Decorator that wraps a test function and provides setup() and teardown()
+    functionality.  The decorated func may define as a parameter any kwargs
+    returned by setup(func_name).  The setup() is func.name aware but not
+    necessarily module aware, so be careful if using the same test function in
+    different test modules
+    """
+    return _with_setup(setup_func, teardown_func, True)(func)
+
+
+@contextmanager
+def inject_into_dag(new_task_dct):
+    """Update (add or replace) tasks in dag with new task config.
+    This should reset any cacheing within the running Stolos instance,
+    but it's not guaranteed.
+    Assumes that the config we're using is the JSONMapping
+    """
+    f = _create_tasks_json(inject=new_task_dct)
+    new_task_dct = JSONMapping(new_task_dct)
+
+    # verify injection worked
+    dg = dt.get_tasks_config()
+    assert isinstance(dg, JSONMapping)
+    dag = dt.build_dag_cached()
+    for k, v in new_task_dct.items():
+        assert dg[k] == v, (
+            "test code: inject_into_dag didn't insert the new tasks?")
+        assert dag.node[k] == dict(v), (
+            "test code: inject_into_dag didn't reset the dag graph")
+
+    yield
+    os.remove(f)
+    dt.build_dag_cached.cache.clear()
+
+
+def enqueue(app_name, job_id, zk, validate_queued=True):
+    # initialize job
+    zkt.maybe_add_subtask(app_name, job_id, zk)
+    zkt.maybe_add_subtask(app_name, job_id, zk)
+    # verify initial conditions
+    if validate_queued:
+        validate_one_queued_task(zk, app_name, job_id)
+
+
+def validate_zero_queued_task(zk, app_name):
+    if zk.exists(join(app_name, 'entries')):
+        nt.assert_equal(
+            0, len(zk.get_children(join(app_name, 'entries'))))
+
+
+def validate_zero_completed_task(zk, app_name):
+    if zk.exists(join(app_name, 'all_subtasks')):
+        nt.assert_equal(
+            0, len(zk.get_children(join(app_name, 'all_subtasks'))))
+
+
+def validate_one_failed_task(zk, app_name, job_id):
+    status = get_zk_status(zk, app_name, job_id)
+    nt.assert_equal(status['num_execute_locks'], 0)
+    nt.assert_equal(status['num_add_locks'], 0)
+    nt.assert_equal(status['in_queue'], False)
+    # nt.assert_equal(status['app_qsize'], 1)
+    nt.assert_equal(status['state'], 'failed')
+
+
+def validate_one_queued_executing_task(zk, app_name, job_id):
+    status = get_zk_status(zk, app_name, job_id)
+    nt.assert_equal(status['num_execute_locks'], 1)
+    nt.assert_equal(status['num_add_locks'], 0)
+    nt.assert_equal(status['in_queue'], True)
+    nt.assert_equal(status['app_qsize'], 1)
+    nt.assert_equal(status['state'], 'pending')
+
+
+def validate_one_queued_task(zk, app_name, job_id):
+    status = get_zk_status(zk, app_name, job_id)
+    nt.assert_equal(status['num_execute_locks'], 0)
+    nt.assert_equal(status['num_add_locks'], 0)
+    nt.assert_equal(status['in_queue'], True)
+    nt.assert_equal(status['app_qsize'], 1)
+    nt.assert_equal(status['state'], 'pending')
+
+
+def validate_one_completed_task(zk, app_name, job_id):
+    status = get_zk_status(zk, app_name, job_id)
+    nt.assert_equal(status['num_execute_locks'], 0)
+    nt.assert_equal(status['num_add_locks'], 0)
+    nt.assert_equal(status['in_queue'], False)
+    nt.assert_equal(status['app_qsize'], 0)
+    nt.assert_equal(status['state'], 'completed')
+
+
+def validate_one_skipped_task(zk, app_name, job_id):
+    status = get_zk_status(zk, app_name, job_id)
+    nt.assert_equal(status['num_execute_locks'], 0)
+    nt.assert_equal(status['num_add_locks'], 0)
+    nt.assert_equal(status['in_queue'], False)
+    nt.assert_equal(status['app_qsize'], 0)
+    nt.assert_equal(status['state'], 'skipped')
+
+
+def cycle_queue(zk, app_name):
+    """Get item from queue, put at back of queue and return item"""
+    q = zk.LockingQueue(app_name)
+    item = q.get()
+    q.put(item)
+    q.consume()
+    return item
+
+
+def consume_queue(zk, app_name, timeout=1):
+    q = zk.LockingQueue(app_name)
+    item = q.get(timeout=timeout)
+    q.consume()
+    return item
+
+
+def get_zk_status(zk, app_name, job_id):
+    path = zkt._get_zookeeper_path(app_name, job_id)
+    elockpath = join(path, 'execute_lock')
+    alockpath = join(path, 'add_lock')
+    entriespath = join(app_name, 'entries')
+    return {
+        'num_add_locks': len(
+            zk.exists(alockpath) and zk.get_children(alockpath) or []),
+        'num_execute_locks': len(
+            zk.exists(elockpath) and zk.get_children(elockpath) or []),
+        'in_queue': (
+            any(zk.get(join(app_name, 'entries', x))[0] == job_id
+                for x in zk.get_children(entriespath))
+            if zk.exists(entriespath) else False),
+        'app_qsize': (
+            len(zk.get_children(entriespath))
+            if zk.exists(entriespath) else 0),
+        'state': zk.get(path)[0],
+    }

stolos/tests/configuration_backend/test_json_config.py

@@ -4,17 +4,18 @@
     JSONMapping, JSONSequence)
 
 
-def with_setup(func):
+def setup_func(func_name):
+    raw = {'a': 1, 'b': [1, 2, 3], 'c': {'aa': 1, 'bb': [1, 2, 3]}}
+    td = JSONMapping(raw)
+    return [], dict(td=td, raw=raw)
+
 
-    def setup_func():
-        raw = {'a': 1, 'b': [1, 2, 3], 'c': {'aa': 1, 'bb': [1, 2, 3]}}
-        td = JSONMapping(raw)
-        return [], dict(td=td, raw=raw)
+def teardown_func():
+    pass
 
-    def teardown_func():
-        pass
 
-    return tt.with_setup(setup_func, teardown_func, True)(func)
+def with_setup(func):
+    return tt.with_setup(func, setup_func, teardown_func)
 
 
 @with_setup

stolos/tests/configuration_backend/test_redis_config.py

@@ -47,8 +47,7 @@ def teardown_func(raw):
 
 
 def with_setup(func):
-    return tt.with_setup(
-        lambda: setup_func(func.func_name), teardown_func, True)(func)
+    return tt.with_setup(func, setup_func, teardown_func)
 
 
 @with_setup