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