tests/data/cli/auto-source-discovery/bt_plugin_test.py

   1 import bt2
   2 import os
   3
   4
   5 class TestIter(bt2._UserMessageIterator):
   6     pass
   7
   8
   9 class Base:
  10     @classmethod
  11     def _print_params(cls, params):
  12         inputs = sorted([str(x) for x in params['inputs']])
  13         print('{}: {}'.format(cls.__name__, ', '.join(inputs)))
  14
  15
  16 @bt2.plugin_component_class
  17 class TestSourceExt(Base, bt2._UserSourceComponent, message_iterator_class=TestIter):
  18     """
  19     Recognize files whose name start with 'aaa', 'bbb' or 'ccc'.
  20
  21     'aaa' files are grouped together, 'bbb' files are grouped together, 'ccc'
  22     files are not grouped.
  23     """
  24
  25     def __init__(self, params):
  26         self._print_params(params)
  27
  28     @staticmethod
  29     def _user_query(query_exec, obj, params, log_level):
  30         if obj == 'support-info':
  31             if params['type'] == 'file':
  32                 name = os.path.basename(str(params['input']))
  33
  34                 if name.startswith('aaa'):
  35                     return {'weight': 1, 'group': 'aaa'}
  36                 elif name.startswith('bbb'):
  37                     return {'weight': 0.5, 'group': 'bbb'}
  38                 elif name.startswith('ccc'):
  39                     # Try two different ways of returning "no group", and two
  40                     # different ways of returning 1 (an int and a float).
  41                     if name[3] == '1':
  42                         return {'weight': 1, 'group': None}
  43                     elif name[3] == '2':
  44                         return {'weight': 1.0, 'group': None}
  45                     elif name[3] == '3':
  46                         return 1
  47                     else:
  48                         return 1.0
  49             else:
  50                 return 0
  51         else:
  52             raise bt2.UnknownObject
  53
  54
  55 @bt2.plugin_component_class
  56 class TestSourceSomeDir(
  57     Base, bt2._UserSourceComponent, message_iterator_class=TestIter
  58 ):
  59     """Recognizes directories named "some-dir".  The file "aaa10" in the
  60     directory "some-dir" won't be found by TestSourceExt, because we won't
  61     recurse in "some-dir"."""
  62
  63     def __init__(self, params):
  64         self._print_params(params)
  65
  66     @staticmethod
  67     def _user_query(query_exec, obj, params, log_level):
  68         if obj == 'support-info':
  69             if params['type'] == 'directory':
  70                 name = os.path.basename(str(params['input']))
  71                 return 1 if name == 'some-dir' else 0
  72             else:
  73                 return 0
  74         else:
  75             raise bt2.UnknownObject
  76
  77
  78 @bt2.plugin_component_class
  79 class TestSourceABCDE(Base, bt2._UserSourceComponent, message_iterator_class=TestIter):
  80     """A source that recognizes the arbitrary string input "ABCDE"."""
  81
  82     def __init__(self, params):
  83         self._print_params(params)
  84
  85     @staticmethod
  86     def _user_query(query_exec, obj, params, log_level):
  87         if obj == 'support-info':
  88             return (
  89                 1.0
  90                 if params['type'] == 'string' and params['input'] == 'ABCDE'
  91                 else 0.0
  92             )
  93         else:
  94             raise bt2.UnknownObject
  95
  96
  97 class TestSourceNoQuery(bt2._UserSourceComponent, message_iterator_class=TestIter):
  98     """A source that does not implement _query at all."""
  99
 100
 101 bt2.register_plugin(module_name=__name__, name="test")