Fixes confluentinc#156 - Adding support for conditional log capture

ducktape/command_line/main.py

@@ -14,6 +14,7 @@
 
 import importlib
 import json
+import yaml
 import os
 import sys
 import traceback
@@ -144,10 +145,16 @@ def main():
     for k, v in args_dict.iteritems():
         session_logger.debug("Configuration: %s=%s", k, v)
 
+    if args_dict['logging_config']:
+        with open(args_dict['logging_config']) as f:
+            logging_config = yaml.load(f.read())
+    else:
+        logging_config = None
+
     # Discover and load tests to be run
     extend_import_paths(args_dict["test_path"])
     loader = TestLoader(session_context, session_logger, repeat=args_dict["repeat"], injected_args=injected_args,
-                        subset=args_dict["subset"], subsets=args_dict["subsets"])
+                        subset=args_dict["subset"], subsets=args_dict["subsets"], logging_config=logging_config)
     try:
         tests = loader.load(args_dict["test_path"])
     except LoaderException as e:

ducktape/command_line/parse_args.py

@@ -25,6 +25,7 @@ def create_ducktape_parser():
     parser = argparse.ArgumentParser(description="Discover and run your tests")
     parser.add_argument('test_path', metavar='test_path', type=str, nargs='*', default=[os.getcwd()],
                         help='one or more space-delimited strings indicating where to search for tests.')
+    parser.add_argument("--logging-config", action="store", help="YAML configuration file for test log collection")
     parser.add_argument("--collect-only", action="store_true", help="display collected tests, but do not run.")
     parser.add_argument("--debug", action="store_true", help="pipe more verbose test output to stdout.")
     parser.add_argument("--config-file", action="store", default=ConsoleDefaults.USER_CONFIG_FILE,

ducktape/tests/condition.py

@@ -0,0 +1,18 @@
+
+class TestCondition(object):
+    def __init__(self, status):
+        self._status = str(status).lower()
+
+    def __eq__(self, other):
+        return str(self).lower() == str(other).lower()
+
+    def __str__(self):
+        return self._status
+
+    def to_json(self):
+        return str(self).upper()
+
+ON_FAIL = TestCondition("on fail")
+ON_PASS = TestCondition("on pass")
+
+

ducktape/tests/loader.py

@@ -44,7 +44,7 @@ class TestLoader(object):
     """Class used to discover and load tests."""
 
     def __init__(self, session_context, logger, repeat=1, injected_args=None, cluster=None, subset=0, subsets=1,
-                 historical_report=None):
+                 historical_report=None, logging_config=None):
         self.session_context = session_context
         self.cluster = cluster
         assert logger is not None
@@ -67,6 +67,9 @@ def __init__(self, session_context, logger, repeat=1, injected_args=None, cluste
         # in any discovered test, whether or not it is parametrized
         self.injected_args = injected_args
 
+
+        self.logging_config = logging_config if logging_config is not None else dict()
+
     def load(self, test_discovery_symbols):
         """Recurse through packages in file hierarchy starting at base_dir, and return a list of test_context objects
         for all discovered tests.
@@ -84,6 +87,7 @@ def load(self, test_discovery_symbols):
         all_test_context_list = []
         for symbol in test_discovery_symbols:
             directory, module_name, cls_name, method_name = self._parse_discovery_symbol(symbol)
+            self.logger.debug("_parse_discovery_symbol %s %s %s %s", directory, module_name, cls_name, method_name)
             directory = os.path.abspath(directory)
 
             test_context_list_for_symbol = self.discover(directory, module_name, cls_name, method_name)
@@ -92,7 +96,11 @@ def load(self, test_discovery_symbols):
             if len(test_context_list_for_symbol) == 0:
                 raise LoaderException("Didn't find any tests for symbol %s." % symbol)
 
+        for test_context in all_test_context_list:
+            self._configure_logging(test_context)
+
         self.logger.debug("Discovered these tests: " + str(all_test_context_list))
+        self.logger.debug("Discovered these tests: " + str(map(type, all_test_context_list)))
 
         # Sort to make sure we get a consistent order for when we create subsets
         all_test_context_list = sorted(all_test_context_list, key=attrgetter("test_id"))
@@ -277,6 +285,67 @@ def _import_modules(self, file_list):
 
         return module_and_file_list
 
+    def _configure_logging(self, test_context):
+        """
+        Configures logging of a service during a test if the service matches the following criteria:
+
+        * The test module name matches the module provided
+        * The test class name matches the class name provided if provided
+        * The test function name matches the function name provided if provided
+        * The service name matches the module name and the class name of the service provide
+            in the format: module_name.class_name
+
+        This function takes a data-structure in logging_config with the following layout:
+
+        {'tests': [{'test': 'test_module_name.TestClassName.test_method_name',
+                    'services': [{
+                                 'name': 'service_module_name.ServiceClassName',
+                                 'logs':
+                                   [{'collect': 'on fail', 'name': 'debug_log'},
+                                    {'collect': True, 'name': 'info_log'},
+                                    {'collect': True, 'name': 'error_log'},
+                                   ]
+                                }]
+                    }]}
+
+        This function will match the test name and the service name with the
+        module, cls, and function in test_context and assign the configured value
+        to log_collect[(log_name, service_name)].
+        """
+        module = test_context.module or test_context.cls.__module__
+        cls = test_context.cls
+        function = test_context.function
+        for test_config in self.logging_config.get('tests', []):
+            test_full_name = test_config.get('test', '')
+            self.logger.debug("test_full_name %s module %s", test_full_name, module)
+            if test_full_name.startswith(module):
+                self.logger.debug("_configure_logging module matches %s", module)
+                rest = test_full_name[len(module):]
+                rest = rest[1:] if rest.startswith(".") else rest
+                self.logger.debug("rest %s cls", rest)
+                if not rest or rest.startswith(cls.__name__):
+                    self.logger.debug("_configure_logging class matches %s", cls)
+                    rest = rest[len(cls.__name__):]
+                    rest = rest[1:] if rest.startswith(".") else rest
+                    self.logger.debug("rest %s cls", rest)
+                    if not rest or rest.startswith(function.__name__):
+                        self.logger.debug("_configure_logging function matches %s", function)
+                        for service in test_config.get('services', []):
+                            for log in service.get('logs', []):
+                                if 'name' in log and 'collect' in log:
+                                    key = (log['name'], service['name'])
+                                    self.logger.debug("_configure_logging (%s) = %s",
+                                                      key,
+                                                      log['collect'])
+                                    if key in test_context.log_collect:
+                                        self.logger.warning("Overwriting log config for %s", key)
+                                    test_context.log_collect[key] = log['collect']
+                                else:
+                                    self.logger.warning("log entry should contain 'name' and 'collect' values")
+        self.logger.debug("log_collect %s", test_context.log_collect)
+
+
+
     def _expand_module(self, module_and_file):
         """Return a list of TestContext objects, one object for every 'testable unit' in module"""
 
@@ -286,6 +355,7 @@ def _expand_module(self, module_and_file):
         module_objects = module.__dict__.values()
         test_classes = [c for c in module_objects if self._is_test_class(c)]
 
+
         for cls in test_classes:
             test_context_list.extend(self._expand_class(
                 TestContext(

ducktape/tests/runner_client.py

@@ -68,13 +68,14 @@ def _sigterm_handler(self, signum, frame):
         """
         os.kill(os.getpid(), signal.SIGINT)
 
-    def _collect_test_context(self, directory, file_name, cls_name, method_name, injected_args):
+    def _collect_test_context(self, directory, file_name, cls_name, method_name, injected_args, log_collect):
         loader = TestLoader(self.session_context, self.logger, injected_args=injected_args, cluster=self.cluster)
         loaded_context_list = loader.discover(directory, file_name, cls_name, method_name)
 
         assert len(loaded_context_list) == 1
         test_context = loaded_context_list[0]
         test_context.cluster = self.cluster
+        test_context.log_collect = log_collect
         return test_context
 
     def run(self):
@@ -122,14 +123,14 @@ def run(self):
 
             data = self.run_test()
 
-            test_status = PASS
+            self.test_context.test_status = test_status = PASS
             self.log(logging.INFO, "PASS")
 
         except BaseException as e:
             err_trace = str(e.message) + "\n" + traceback.format_exc(limit=16)
             self.log(logging.INFO, "FAIL: " + err_trace)
 
-            test_status = FAIL
+            self.test_context.test_status = test_status = FAIL
             summary += err_trace
 
         finally:

ducktape/tests/test.py

@@ -26,9 +26,12 @@
 from ducktape.services.service_registry import ServiceRegistry
 from ducktape.template import TemplateRenderer
 from ducktape.mark.resource import CLUSTER_SIZE_KEYWORD
+from ducktape.tests.status import FAIL
+from ducktape.tests.condition import ON_FAIL
 
 
 class Test(TemplateRenderer):
+
     """Base class for tests.
     """
     def __init__(self, test_context, *args, **kwargs):
@@ -109,7 +112,7 @@ def copy_service_logs(self):
         for service in self.test_context.services:
             if not hasattr(service, 'logs') or len(service.logs) == 0:
                 self.test_context.logger.debug("Won't collect service logs from %s - no logs to collect." %
-                    service.service_id)
+                                               service.service_id)
                 continue
 
             log_dirs = service.logs
@@ -145,6 +148,14 @@ def copy_service_logs(self):
                              'service': service,
                              'message': e.message})
 
+    def mark_for_collect_on_failure(self, service, log_name=None):
+        if log_name is None:
+            # Mark every log for collection
+            for log_name in service.logs:
+                self.test_context.log_collect[(log_name, service)] = ON_FAIL
+        else:
+            self.test_context.log_collect[(log_name, service)] = ON_FAIL
+
     def mark_for_collect(self, service, log_name=None):
         if log_name is None:
             # Mark every log for collection
@@ -158,9 +169,14 @@ def mark_no_collect(self, service, log_name=None):
 
     def should_collect_log(self, log_name, service):
         key = (log_name, service)
+        service_name = ".".join([service.__class__.__module__,
+                                 service.__class__.__name__])
         default = service.logs[log_name]["collect_default"]
         val = self.test_context.log_collect.get(key, default)
-        return val
+        val = self.test_context.log_collect.get((log_name, service_name), val)
+        if val == ON_FAIL and self.test_context.test_status == FAIL:
+            return True
+        return val is True
 
 
 def _compress_cmd(log_path):
@@ -276,6 +292,7 @@ def __init__(self, **kwargs):
 
         self._logger = None
         self._local_scratch_dir = None
+        self.test_status = None
 
     def __repr__(self):
         return "<module=%s, cls=%s, function=%s, injected_args=%s, file=%s, ignore=%s, cluster_size=%s>" % \
@@ -305,7 +322,8 @@ def test_metadata(self):
             "file_name": os.path.basename(self.file),
             "cls_name": self.cls.__name__,
             "method_name": self.function.__name__,
-            "injected_args": self.injected_args
+            "injected_args": self.injected_args,
+            "log_collect": self.log_collect
         }
 
     @staticmethod

setup.py

@@ -42,7 +42,7 @@ def run_tests(self):
       url="http://github.com/confluentinc/ducktape",
       packages=find_packages(),
       package_data={'ducktape': ['templates/report/*']},
-      install_requires=['jinja2', 'requests', 'paramiko', 'pysistence', 'pyzmq'],
+      install_requires=['jinja2', 'requests', 'paramiko', 'pysistence', 'pyzmq', 'pyyaml'],
       tests_require=['pytest', 'mock', 'psutil==4.1.0', 'memory_profiler==0.41', 'statistics', 'requests-testadapter'],
       cmdclass={'test': PyTest},
       )

tests/command_line/check_parse_args.py

@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 from ducktape.command_line.parse_args import parse_args
 
+from ..loader.check_loader import logging_config_directory
+
 from cStringIO import StringIO
 from exceptions import SystemExit
 
@@ -139,3 +142,11 @@ def check_config_file_option(self):
             assert args_dict["parameters"] == "PARAMETERS-user"
         finally:
             shutil.rmtree(tmpdir)
+
+    def check_logging_config_file_option(self):
+        """Check that the logging config file option works"""
+
+        logging_config_file = os.path.join(logging_config_directory(), 'test_a.yaml')
+
+        args_dict = parse_args(["--logging-config", logging_config_file])
+        assert args_dict["logging_config"] == logging_config_file