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

   1 import bt2
   2 import os
   3
   4 # Source component classes in this file recognize and group inputs in
   5 # various ways.  One stream is created by each component, with a name
   6 # derived from the component class and the inputs.  This is then checked
   7 # using the `sink.text.details` sink.
   8
   9
  10 class TestIter(bt2._UserMessageIterator):
  11     def __init__(self, output_port):
  12         inputs = output_port.user_data['inputs']
  13         sc = output_port.user_data['sc']
  14         tc = sc.trace_class
  15         t = tc()
  16         s = t.create_stream(sc, name=self._make_stream_name(inputs))
  17
  18         self._msgs = [
  19             self._create_stream_beginning_message(s),
  20             self._create_stream_end_message(s),
  21         ]
  22
  23     def _make_stream_name(self, inputs):
  24         comp_cls_name = self._component.__class__.__name__
  25         return (
  26             comp_cls_name
  27             + ': '
  28             + ', '.join(sorted([os.path.basename(str(x)) for x in inputs]))
  29         )
  30
  31     def __next__(self):
  32         if len(self._msgs) == 0:
  33             raise StopIteration
  34
  35         return self._msgs.pop(0)
  36
  37
  38 class Base:
  39     def __init__(self, params):
  40         tc = self._create_trace_class()
  41         sc = tc.create_stream_class()
  42
  43         self._add_output_port('out', {'inputs': params['inputs'], 'sc': sc})
  44
  45
  46 @bt2.plugin_component_class
  47 class TestSourceExt(Base, bt2._UserSourceComponent, message_iterator_class=TestIter):
  48     """
  49     Recognize files whose name start with 'aaa', 'bbb' or 'ccc'.
  50
  51     'aaa' files are grouped together, 'bbb' files are grouped together, 'ccc'
  52     files are not grouped.
  53     """
  54
  55     def __init__(self, config, params, obj):
  56         super().__init__(params)
  57
  58     @staticmethod
  59     def _user_query(priv_query_exec, obj, params, method_obj):
  60         if obj == 'babeltrace.support-info':
  61             if params['type'] == 'file':
  62                 name = os.path.basename(str(params['input']))
  63
  64                 if name.startswith('aaa'):
  65                     return {'weight': 1, 'group': 'aaa'}
  66                 elif name.startswith('bbb'):
  67                     return {'weight': 0.5, 'group': 'bbb'}
  68                 elif name.startswith('ccc'):
  69                     # Try two different ways of returning 1 (an int and a float).
  70                     if name[3] == '1':
  71                         return 1
  72                     else:
  73                         return 1.0
  74
  75             return 0
  76         else:
  77             raise bt2.UnknownObject
  78
  79
  80 @bt2.plugin_component_class
  81 class TestSourceSomeDir(
  82     Base, bt2._UserSourceComponent, message_iterator_class=TestIter
  83 ):
  84     """Recognizes directories named "some-dir".  The file "aaa10" in the
  85     directory "some-dir" won't be found by TestSourceExt, because we won't
  86     recurse in "some-dir"."""
  87
  88     def __init__(self, config, params, obj):
  89         super().__init__(params)
  90
  91     @staticmethod
  92     def _user_query(priv_query_exec, obj, params, method_obj):
  93         if obj == 'babeltrace.support-info':
  94             if params['type'] == 'directory':
  95                 name = os.path.basename(str(params['input']))
  96                 return 1 if name == 'some-dir' else 0
  97             else:
  98                 return 0
  99         else:
 100             raise bt2.UnknownObject
 101
 102
 103 @bt2.plugin_component_class
 104 class TestSourceABCDE(Base, bt2._UserSourceComponent, message_iterator_class=TestIter):
 105     """A source that recognizes the arbitrary string input "ABCDE"."""
 106
 107     def __init__(self, config, params, obj):
 108         super().__init__(params)
 109
 110     @staticmethod
 111     def _user_query(priv_query_exec, obj, params, method_obj):
 112         if obj == 'babeltrace.support-info':
 113             return (
 114                 1.0
 115                 if params['type'] == 'string' and params['input'] == 'ABCDE'
 116                 else 0.0
 117             )
 118         else:
 119             raise bt2.UnknownObject
 120
 121
 122 class TestSourceNoQuery(bt2._UserSourceComponent, message_iterator_class=TestIter):
 123     """A source that does not implement _query at all."""
 124
 125
 126 bt2.register_plugin(module_name=__name__, name="test")