18c396eb423d2498d97cebbfdd0a68fce0b576b0
[babeltrace.git] / 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, config, 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")
This page took 0.031443 seconds and 3 git commands to generate.