Coverage for auttcomp/testing/test_extensions.py: 100%
142 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-24 12:00 -0600
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-24 12:00 -0600
1from ..extensions import KeyValuePair, Api as f
2from .base_test import get_hugging_face_sample, get_civitai_sample
3from ..quicklog import tracelog
5@tracelog("test_id")
6def test_id():
7 func = f.id(123)
8 actual = func()
9 assert actual == 123
11@tracelog("test_at")
12def test_at():
13 data = {"foo": 123}
14 actual = f.at(lambda x: x['foo'])(data)
15 assert actual == 123
17@tracelog("test_map")
18def test_map():
19 data = [1, 2, 3]
20 gen = f.map(lambda x: x + 1)(data)
21 actual = list(gen)
22 assert actual == [2, 3, 4]
24@tracelog("test_filter")
25def test_filter():
26 data = [1, 2, 3]
27 gen = f.filter(lambda x: x % 2 == 0)(data)
28 actual = list(gen)
29 assert actual == [2]
31@tracelog("test_reduce")
32def test_reduce():
33 data = [2, 2, 2]
34 actual = f.reduce(lambda p, n: p + n)(data)
35 assert actual == 6
37@tracelog("test_flatmap")
38def test_flatmap():
39 data = [[1], [1], [1]]
40 gen = f.flatmap(lambda x: x)(data)
41 actual = list(gen)
42 assert actual == [1, 1, 1]
44@tracelog("test_flatmap_id")
45def test_flatmap_id():
46 data = [[1], [1], [1]]
47 gen = f.flatmap()(data)
48 actual = list(gen)
49 assert actual == [1, 1, 1]
51@tracelog("test_reverse")
52def test_reverse():
53 data = [1, 2, 3]
54 gen = f.reverse(data)
55 actual = list(gen)
56 assert actual == [3, 2, 1]
58@tracelog("test_any")
59def test_any():
60 data1 = [0, 0, 0]
61 data2 = [0, 0, 111]
63 actual1 = f.any(lambda x: x == 111)(data1)
64 actual2 = f.any(lambda x: x == 111)(data2)
66 assert actual1 == False
67 assert actual2 == True
69@tracelog("test_all")
70def test_all():
71 data1 = [0, 0, 0]
72 data2 = [0, 0, 111]
74 actual1 = f.all(lambda x: x == 0)(data1)
75 actual2 = f.all(lambda x: x == 0)(data2)
77 assert actual1 == True
78 assert actual2 == False
80@tracelog("test_sort")
81def test_sort():
82 data = [2, 3, 1]
83 gen = f.sort(data)
84 actual = list(gen)
85 assert actual == [1, 2, 3]
87@tracelog("test_sort_by")
88def test_sort_by():
89 data = [
90 {"id": 2},
91 {"id": 3},
92 {"id": 1}
93 ]
94 expected = [
95 {"id": 1},
96 {"id": 2},
97 {"id": 3}
98 ]
100 gen = f.sort_by(lambda x: x['id'])(data)
102 actual = list(gen)
104 assert actual == expected
106@tracelog("test_sort_by_desc")
107def test_sort_by_desc():
108 data = [
109 {"id": 2},
110 {"id": 3},
111 {"id": 1}
112 ]
113 expected = [
114 {"id": 3},
115 {"id": 2},
116 {"id": 1}
117 ]
119 gen = f.sort_by_desc(lambda x: x['id'])(data)
121 actual = list(gen)
123 assert actual == expected
125@tracelog("test_take")
126def test_take():
127 data = [1, 2, 3]
128 gen = f.take(2)(data)
129 actual = list(gen)
130 assert actual == [1, 2]
132@tracelog("test_skip")
133def test_skip():
134 data = [1, 2, 3]
135 gen = f.skip(1)(data)
136 actual = list(gen)
137 assert actual == [2, 3]
139@tracelog("test_group")
140def test_group():
141 data = [
142 {"id": 1, "tag": "TAG1"},
143 {"id": 2, "tag": "TAG2"},
144 {"id": 3, "tag": "TAG1"},
145 {"id": 4, "tag": "TAG2"}
146 ]
148 expected = [
149 KeyValuePair("TAG1", [
150 {"id": 1, "tag": "TAG1"},
151 {"id": 3, "tag": "TAG1"}
152 ]),
153 KeyValuePair("TAG2", [
154 {"id": 2, "tag": "TAG2"},
155 {"id": 4, "tag": "TAG2"}
156 ])
157 ]
159 gen = f.group(lambda x: x['tag'])
161 actual = list(gen(data))
163 assert actual == expected
166@tracelog("test_join")
167def test_join():
169 dataFoo = [
170 {"foo_id": 1, "foo": "foo1"},
171 {"foo_id": 2, "foo": "foo2"},
172 {"foo_id": 50, "foo": "foo50"},
173 {"foo_id": 3, "foo": "foo3"},
174 {"foo_id": 4, "foo": "foo4"}
175 ]
177 dataBar = [
178 {"bar_id": 1, "bar": "bar1"},
179 {"bar_id": 2, "bar": "bar2"},
180 {"bar_id": 3, "bar": "bar3"},
181 {"bar_id": 4, "bar": "bar4"},
182 {"bar_id": 100, "bar": "bar100"}
183 ]
185 expected = [
186 (1, (["foo1"], ["bar1"])),
187 (2, (["foo2"], ["bar2"])),
188 (3, (["foo3"], ["bar3"])),
189 (4, (["foo4"], ["bar4"]))
190 ]
192 gen = f.join(
193 left_data=dataFoo,
194 left_key_func=lambda x: x['foo_id'],
195 right_key_func=lambda x: x['bar_id'],
196 left_value_selector=lambda x: x['foo'],
197 right_value_selector=lambda x: x['bar']
198 )(dataBar)
200 actual = list(gen)
202 assert actual == expected
204@tracelog("test_distinct_set")
205def test_distinct_set():
206 arr = f.id([1, 2, 3, 3, 3, 3])
207 slow_distinct = arr > f.distinct
208 fast_distinct = arr > f.distinct_set
209 assert slow_distinct == fast_distinct
212@tracelog("test_zip")
213def test_zip():
214 data1 = ["a", "b", "c"]
215 data2 = ["A", "B"]
216 expected = [("a", "A"),("b", "B"),("c", None)]
218 gen = f.zip(data2)(data1)
220 actual = list(gen)
222 assert actual == expected
224@tracelog("test_flatnest")
225def test_flatnest():
226 data = {
227 "depth": 1,
228 "nest": {
229 "depth": 2,
230 "nest": {
231 "depth": 3,
232 "nest": None
233 }
234 }
235 }
237 expected = [1, 2, 3]
239 gen = f.flatnest(
240 path_selector=lambda x: x['nest'],
241 data_selector=lambda x: x['depth']
242 )(data)
244 actual = list(gen)
246 assert actual == expected
248@tracelog("test_huggingface_sample")
249def test_huggingface_sample():
250 data = get_hugging_face_sample()
252 result = f.id(data.models) > (
253 f.group(lambda x: x.author)
254 | f.map(lambda x: (
255 x.key,
256 f.id(x.value) > f.map(lambda x2: x2.downloads) | sum
257 ))
258 | f.sort_by_desc(lambda x: x[1])
259 | f.take(3)
260 | list
261 )
263 assert result == [('black-forest-labs', 1548084), ('deepseek-ai', 1448374), ('microsoft', 264891)]
265@tracelog("test_civitai_sample")
266def test_civitai_sample():
267 data = get_civitai_sample()
269 most_common_tags = f.id(data.result.data.json.collection.items) > (
270 f.flatmap(lambda x: x.data.tagIds)
271 | f.group()
272 | f.map(lambda x: (
273 x.key,
274 f.id(x.value) > f(len)
275 ))
276 | f.sort_by_desc(lambda x: x[1])
277 | f.map(lambda x: x[0])
278 | f.take(3)
279 | list
280 )
282 assert most_common_tags == [292, 81, 5262]