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

1from ..extensions import KeyValuePair, Api as f 

2from .base_test import get_hugging_face_sample, get_civitai_sample 

3from ..quicklog import tracelog 

4 

5@tracelog("test_id") 

6def test_id(): 

7 func = f.id(123) 

8 actual = func() 

9 assert actual == 123 

10 

11@tracelog("test_at") 

12def test_at(): 

13 data = {"foo": 123} 

14 actual = f.at(lambda x: x['foo'])(data) 

15 assert actual == 123 

16 

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] 

23 

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] 

30 

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 

36 

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] 

43 

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] 

50 

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] 

57 

58@tracelog("test_any") 

59def test_any(): 

60 data1 = [0, 0, 0] 

61 data2 = [0, 0, 111] 

62 

63 actual1 = f.any(lambda x: x == 111)(data1) 

64 actual2 = f.any(lambda x: x == 111)(data2) 

65 

66 assert actual1 == False 

67 assert actual2 == True 

68 

69@tracelog("test_all") 

70def test_all(): 

71 data1 = [0, 0, 0] 

72 data2 = [0, 0, 111] 

73 

74 actual1 = f.all(lambda x: x == 0)(data1) 

75 actual2 = f.all(lambda x: x == 0)(data2) 

76 

77 assert actual1 == True 

78 assert actual2 == False 

79 

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] 

86 

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 ] 

99 

100 gen = f.sort_by(lambda x: x['id'])(data) 

101 

102 actual = list(gen) 

103 

104 assert actual == expected 

105 

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 ] 

118 

119 gen = f.sort_by_desc(lambda x: x['id'])(data) 

120 

121 actual = list(gen) 

122 

123 assert actual == expected 

124 

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] 

131 

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] 

138 

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 ] 

147 

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 ] 

158 

159 gen = f.group(lambda x: x['tag']) 

160 

161 actual = list(gen(data)) 

162 

163 assert actual == expected 

164 

165 

166@tracelog("test_join") 

167def test_join(): 

168 

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 ] 

176 

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 ] 

184 

185 expected = [ 

186 (1, (["foo1"], ["bar1"])), 

187 (2, (["foo2"], ["bar2"])), 

188 (3, (["foo3"], ["bar3"])), 

189 (4, (["foo4"], ["bar4"])) 

190 ] 

191 

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) 

199 

200 actual = list(gen) 

201 

202 assert actual == expected 

203 

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 

210 

211 

212@tracelog("test_zip") 

213def test_zip(): 

214 data1 = ["a", "b", "c"] 

215 data2 = ["A", "B"] 

216 expected = [("a", "A"),("b", "B"),("c", None)] 

217 

218 gen = f.zip(data2)(data1) 

219 

220 actual = list(gen) 

221 

222 assert actual == expected 

223 

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 } 

236 

237 expected = [1, 2, 3] 

238 

239 gen = f.flatnest( 

240 path_selector=lambda x: x['nest'], 

241 data_selector=lambda x: x['depth'] 

242 )(data) 

243 

244 actual = list(gen) 

245 

246 assert actual == expected 

247 

248@tracelog("test_huggingface_sample") 

249def test_huggingface_sample(): 

250 data = get_hugging_face_sample() 

251 

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 ) 

262 

263 assert result == [('black-forest-labs', 1548084), ('deepseek-ai', 1448374), ('microsoft', 264891)] 

264 

265@tracelog("test_civitai_sample") 

266def test_civitai_sample(): 

267 data = get_civitai_sample() 

268 

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 ) 

281 

282 assert most_common_tags == [292, 81, 5262] 

283