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

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