tests/data/auto-source-discovery/params-log-level/bt_plugin_test.py

   1 # SPDX-License-Identifier: GPL-2.0-only
   2 #
   3 # Copyright (C) 2019 EfficiOS Inc.
   4 #
   5
   6 import os
   7
   8 import bt2
   9
  10 # This file defines source component classes to help verify the parameters an
  11 # log levels passed to components.  Each component creates one stream, with a
  12 # name derived from either:
  13 #
  14 #   - the received params that start with `test-`
  15 #   - the received log level
  16 #
  17 # The `what` parameter determines what is used.
  18
  19
  20 class TestIter(bt2._UserMessageIterator):
  21     def __init__(self, config, output_port):
  22         params = output_port.user_data["params"]
  23         obj = output_port.user_data["obj"]
  24
  25         comp_cls_name = self._component.__class__.__name__
  26
  27         if params["what"] == "test-params":
  28             items = sorted([str(x) for x in params.items() if x[0].startswith("test-")])
  29             stream_name = "{}: {}".format(comp_cls_name, ", ".join(items))
  30         elif params["what"] == "log-level":
  31             log_level = self._component.logging_level
  32             stream_name = "{}: {}".format(comp_cls_name, log_level)
  33         elif params["what"] == "python-obj":
  34             assert type(obj) is str or obj is None
  35             stream_name = "{}: {}".format(comp_cls_name, obj)
  36         else:
  37             assert False
  38
  39         sc = output_port.user_data["sc"]
  40         tc = sc.trace_class
  41         t = tc()
  42         s = t.create_stream(sc, name=stream_name)
  43
  44         self._msgs = [
  45             self._create_stream_beginning_message(s),
  46             self._create_stream_end_message(s),
  47         ]
  48
  49     def __next__(self):
  50         if len(self._msgs) == 0:
  51             raise StopIteration
  52
  53         return self._msgs.pop(0)
  54
  55
  56 class Base:
  57     def __init__(self, params, obj):
  58         tc = self._create_trace_class()
  59         sc = tc.create_stream_class()
  60
  61         self._add_output_port("out", {"params": params, "obj": obj, "sc": sc})
  62
  63
  64 @bt2.plugin_component_class
  65 class TestSourceA(Base, bt2._UserSourceComponent, message_iterator_class=TestIter):
  66     def __init__(self, config, params, obj):
  67         super().__init__(params, obj)
  68
  69     @staticmethod
  70     def _user_query(priv_query_exec, obj, params, method_obj):
  71         # Match files starting with 'aaa'.
  72
  73         if obj == "babeltrace.support-info":
  74             if params["type"] != "file":
  75                 return 0
  76
  77             name = os.path.basename(str(params["input"]))
  78
  79             if name.startswith("aaa"):
  80                 return {"weight": 1, "group": "aaa"}
  81             else:
  82                 return 0
  83         else:
  84             raise bt2.UnknownObject
  85
  86
  87 @bt2.plugin_component_class
  88 class TestSourceB(Base, bt2._UserSourceComponent, message_iterator_class=TestIter):
  89     def __init__(self, config, params, obj):
  90         super().__init__(params, obj)
  91
  92     @staticmethod
  93     def _user_query(priv_query_exec, obj, params, method_obj):
  94         # Match files starting with 'bbb'.
  95
  96         if obj == "babeltrace.support-info":
  97             if params["type"] != "file":
  98                 return 0
  99
 100             name = os.path.basename(str(params["input"]))
 101
 102             if name.startswith("bbb"):
 103                 return {"weight": 1, "group": "bbb"}
 104             else:
 105                 return 0
 106         else:
 107             raise bt2.UnknownObject
 108
 109
 110 bt2.register_plugin(module_name=__name__, name="test")