Skip to content

Commit aca99a5

Browse files
petrmitrichevcopybara-github
petrmitrichev
authored and
copybara-github
committed
Use test::DataSlice and test::DataItem to make functor tests less verbose.
PiperOrigin-RevId: 720584921 Change-Id: I72de75aaca9768aea84ee072708422709537dc51
1 parent b09ff27 commit aca99a5

File tree

5 files changed

+32
-114
lines changed

5 files changed

+32
-114
lines changed

‎koladata/functor/BUILD‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -247,8 +247,6 @@ cc_test(
247247
"//koladata:test_utils",
248248
"//koladata/expr:expr_eval",
249249
"//koladata/internal:data_item",
250-
"//koladata/internal:data_slice",
251-
"//koladata/internal:dtype",
252250
"//koladata/internal/testing:matchers",
253251
"//koladata/testing:matchers",
254252
"//koladata/testing:test_env",
@@ -260,7 +258,6 @@ cc_test(
260258
"@com_google_arolla//arolla/dense_array",
261259
"@com_google_arolla//arolla/expr",
262260
"@com_google_arolla//arolla/jagged_shape/testing",
263-
"@com_google_arolla//arolla/qtype",
264261
"@com_google_arolla//arolla/util",
265262
"@com_google_arolla//arolla/util:status_backport",
266263
"@com_google_googletest//:gtest_main",
@@ -289,6 +286,7 @@ cc_test(
289286
":signature",
290287
":signature_storage",
291288
"//koladata:data_slice",
289+
"//koladata:test_utils",
292290
"//koladata/expr:expr_eval",
293291
"//koladata/internal:data_item",
294292
"//koladata/internal:dtype",
@@ -320,7 +318,6 @@ cc_test(
320318
"//koladata:data_slice",
321319
"//koladata:object_factories",
322320
"//koladata:test_utils",
323-
"//koladata/internal:data_item",
324321
"//koladata/internal:dtype",
325322
"//koladata/operators:lib",
326323
"//koladata/s11n",

‎koladata/functor/call_test.cc‎

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,7 @@
3232
#include "koladata/functor/functor.h"
3333
#include "koladata/functor/signature.h"
3434
#include "koladata/functor/signature_storage.h"
35-
#include "koladata/internal/data_item.h"
36-
#include "koladata/internal/dtype.h"
35+
#include "koladata/test_utils.h"
3736
#include "koladata/testing/matchers.h"
3837
#include "arolla/expr/expr.h"
3938
#include "arolla/expr/expr_node.h"
@@ -68,9 +67,7 @@ absl::StatusOr<arolla::expr::ExprNodePtr> CreateVariable(
6867
absl::StatusOr<DataSlice> WrapExpr(
6968
absl::StatusOr<arolla::expr::ExprNodePtr> expr_or_error) {
7069
ASSIGN_OR_RETURN(auto expr, expr_or_error);
71-
return DataSlice::Create(
72-
internal::DataItem(arolla::expr::ExprQuote(std::move(expr))),
73-
internal::DataItem(schema::kExpr));
70+
return test::DataItem(arolla::expr::ExprQuote(std::move(expr)));
7471
}
7572

7673
TEST(CallTest, VariableRhombus) {
@@ -196,9 +193,7 @@ TEST(CallTest, DataSliceVariable) {
196193
ASSERT_OK_AND_ASSIGN(auto koda_signature,
197194
CppSignatureToKodaSignature(signature));
198195
ASSERT_OK_AND_ASSIGN(auto returns_expr, WrapExpr(CreateVariable("a")));
199-
ASSERT_OK_AND_ASSIGN(auto var_a,
200-
DataSlice::Create(internal::DataItem(57),
201-
internal::DataItem(schema::kInt32)));
196+
auto var_a = test::DataItem(57);
202197
ASSERT_OK_AND_ASSIGN(
203198
auto fn, CreateFunctor(returns_expr, koda_signature, {{"a", var_a}}));
204199
ASSERT_OK_AND_ASSIGN(auto result, CallFunctorWithCompilationCache(
@@ -223,9 +218,7 @@ TEST(CallTest, EvalError) {
223218
"math.add", {CreateInput("a"), arolla::expr::Literal(57)})));
224219
ASSERT_OK_AND_ASSIGN(auto fn, CreateFunctor(returns_expr, koda_signature,
225220
{{"foo", var_expr}}));
226-
ASSERT_OK_AND_ASSIGN(auto input,
227-
DataSlice::Create(internal::DataItem(43),
228-
internal::DataItem(schema::kInt32)));
221+
auto input = test::DataItem(43);
229222
// This error message should be improved, in particular it should actually
230223
// mention that we are evaluating a functor, which variable, etc.
231224
// It is OK to only improve this on the Python side, the C++ error is not

‎koladata/functor/functor_test.cc‎

Lines changed: 8 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -69,9 +69,7 @@ absl::StatusOr<arolla::expr::ExprNodePtr> CreateVariable(
6969
absl::StatusOr<DataSlice> WrapExpr(
7070
absl::StatusOr<arolla::expr::ExprNodePtr> expr_or_error) {
7171
ASSIGN_OR_RETURN(auto expr, expr_or_error);
72-
return DataSlice::Create(
73-
internal::DataItem(arolla::expr::ExprQuote(std::move(expr))),
74-
internal::DataItem(schema::kExpr));
72+
return test::DataItem(arolla::expr::ExprQuote(std::move(expr)));
7573
}
7674

7775
TEST(CreateFunctorTest, Basic) {
@@ -87,9 +85,7 @@ TEST(CreateFunctorTest, Basic) {
8785
WrapExpr(arolla::expr::CallOp("math.multiply",
8886
{CreateInput("a"), CreateVariable("a")})));
8987
ASSERT_OK_AND_ASSIGN(auto var_a_expr, WrapExpr(CreateInput("b")));
90-
ASSERT_OK_AND_ASSIGN(auto slice_57,
91-
DataSlice::Create(internal::DataItem(57),
92-
internal::DataItem(schema::kInt32)));
88+
auto slice_57 = test::DataItem(57);
9389
ASSERT_OK_AND_ASSIGN(auto my_obj, ObjectCreator::FromAttrs(
9490
DataBag::Empty(), {"a"}, {slice_57}));
9591
ASSERT_OK_AND_ASSIGN(auto fn,
@@ -118,9 +114,7 @@ TEST(CreateFunctorTest, DefaultSignature) {
118114
WrapExpr(arolla::expr::CallOp("math.multiply",
119115
{CreateInput("a"), CreateVariable("a")})));
120116
ASSERT_OK_AND_ASSIGN(auto var_a_expr, WrapExpr(CreateInput("b")));
121-
ASSERT_OK_AND_ASSIGN(auto slice_57,
122-
DataSlice::Create(internal::DataItem(57),
123-
internal::DataItem(schema::kInt32)));
117+
auto slice_57 = test::DataItem(57);
124118
ASSERT_OK_AND_ASSIGN(
125119
auto fn, CreateFunctor(returns_expr, std::nullopt, {{"a", var_a_expr}}));
126120
ASSERT_OK_AND_ASSIGN(auto koda_signature, fn.GetAttr(kSignatureAttrName));
@@ -139,9 +133,7 @@ TEST(CreateFunctorTest, DefaultSignature) {
139133
}
140134

141135
TEST(CreateFunctorTest, NonExprReturns) {
142-
ASSERT_OK_AND_ASSIGN(auto slice_57,
143-
DataSlice::Create(internal::DataItem(57),
144-
internal::DataItem(schema::kInt32)));
136+
auto slice_57 = test::DataItem(57);
145137
ASSERT_OK_AND_ASSIGN(auto signature, Signature::Create({}));
146138
ASSERT_OK_AND_ASSIGN(auto koda_signature,
147139
CppSignatureToKodaSignature(signature));
@@ -151,9 +143,7 @@ TEST(CreateFunctorTest, NonExprReturns) {
151143
}
152144

153145
TEST(CreateFunctorTest, Non0RankReturns) {
154-
ASSERT_OK_AND_ASSIGN(auto slice_57,
155-
DataSlice::Create(internal::DataItem(57),
156-
internal::DataItem(schema::kInt32)));
146+
auto slice_57 = test::DataItem(57);
157147
ASSERT_OK_AND_ASSIGN(
158148
auto slice_57_1dim,
159149
slice_57.Reshape(DataSlice::JaggedShape::FlatFromSize(1)));
@@ -168,9 +158,7 @@ TEST(CreateFunctorTest, Non0RankReturns) {
168158

169159
TEST(CreateFunctorTest, InvalidSignature) {
170160
ASSERT_OK_AND_ASSIGN(auto returns_expr, WrapExpr(arolla::expr::Literal(57)));
171-
ASSERT_OK_AND_ASSIGN(auto slice_57,
172-
DataSlice::Create(internal::DataItem(57),
173-
internal::DataItem(schema::kInt32)));
161+
auto slice_57 = test::DataItem(57);
174162
EXPECT_THAT(
175163
CreateFunctor(returns_expr, slice_57, {}),
176164
StatusIs(absl::StatusCode::kInvalidArgument,
@@ -182,9 +170,7 @@ TEST(CreateFunctorTest, VariablesWithNon0Rank) {
182170
ASSERT_OK_AND_ASSIGN(auto signature, Signature::Create({}));
183171
ASSERT_OK_AND_ASSIGN(auto koda_signature,
184172
CppSignatureToKodaSignature(signature));
185-
ASSERT_OK_AND_ASSIGN(auto slice_57,
186-
DataSlice::Create(internal::DataItem(57),
187-
internal::DataItem(schema::kInt32)));
173+
auto slice_57 = test::DataItem(57);
188174
ASSERT_OK_AND_ASSIGN(
189175
auto slice_57_with_dim,
190176
slice_57.Reshape(DataSlice::JaggedShape::FlatFromSize(1)));
@@ -213,7 +199,7 @@ TEST(IsFunctorTest, Basic) {
213199
}
214200

215201
TEST(IsFunctorTest, PrimitiveWithBag) {
216-
auto x = test::DataItem<int32_t>(1).WithBag(DataBag::Empty());
202+
auto x = test::DataItem(1).WithBag(DataBag::Empty());
217203
EXPECT_THAT(IsFunctor(x), IsOkAndHolds(false));
218204
EXPECT_THAT(IsFunctor(*x.WithSchema(test::Schema(schema::kObject))),
219205
IsOkAndHolds(false));

‎koladata/functor/map_test.cc‎

Lines changed: 13 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,6 @@
3636
#include "koladata/functor/signature.h"
3737
#include "koladata/functor/signature_storage.h"
3838
#include "koladata/internal/data_item.h"
39-
#include "koladata/internal/data_slice.h"
40-
#include "koladata/internal/dtype.h"
4139
#include "koladata/internal/testing/matchers.h"
4240
#include "koladata/test_utils.h"
4341
#include "koladata/testing/matchers.h"
@@ -47,7 +45,6 @@
4745
#include "arolla/expr/expr_node.h"
4846
#include "arolla/expr/quote.h"
4947
#include "arolla/jagged_shape/testing/matchers.h"
50-
#include "arolla/qtype/typed_ref.h"
5148
#include "arolla/util/cancellation_context.h"
5249
#include "arolla/util/text.h"
5350
#include "arolla/util/status_macros_backport.h"
@@ -69,9 +66,7 @@ absl::StatusOr<arolla::expr::ExprNodePtr> CreateInput(absl::string_view name) {
6966
absl::StatusOr<DataSlice> WrapExpr(
7067
absl::StatusOr<arolla::expr::ExprNodePtr> expr_or_error) {
7168
ASSIGN_OR_RETURN(auto expr, expr_or_error);
72-
return DataSlice::Create(
73-
internal::DataItem(arolla::expr::ExprQuote(std::move(expr))),
74-
internal::DataItem(schema::kExpr));
69+
return test::DataItem(arolla::expr::ExprQuote(std::move(expr)));
7570
}
7671

7772
TEST(MapTest, Basic) {
@@ -93,31 +88,12 @@ TEST(MapTest, Basic) {
9388
ASSERT_OK_AND_ASSIGN(auto fn_b,
9489
CreateFunctor(returns_b_expr, koda_signature, {}));
9590
auto merged_bag = DataBag::CommonDataBag({fn_a.GetBag(), fn_b.GetBag()});
96-
ASSERT_OK_AND_ASSIGN(
97-
DataSlice fn,
98-
DataSlice::CreateWithFlatShape(
99-
internal::DataSliceImpl::Create(
100-
{fn_a.item(), fn_b.item(), internal::DataItem(), fn_a.item()}),
101-
fn_a.GetSchemaImpl(), merged_bag));
102-
ASSERT_OK_AND_ASSIGN(
103-
DataSlice a_input,
104-
DataSlice::CreateWithFlatShape(
105-
internal::DataSliceImpl::Create(
106-
arolla::CreateFullDenseArray<int32_t>({1, 2, 3, 4})),
107-
internal::DataItem(schema::kInt32)));
108-
ASSERT_OK_AND_ASSIGN(
109-
DataSlice b_input,
110-
DataSlice::CreateWithFlatShape(
111-
internal::DataSliceImpl::Create(
112-
arolla::CreateDenseArray<int32_t>({5, 6, 7, std::nullopt})),
113-
internal::DataItem(schema::kInt32)));
114-
115-
ASSERT_OK_AND_ASSIGN(
116-
DataSlice expected1,
117-
DataSlice::CreateWithFlatShape(
118-
internal::DataSliceImpl::Create(arolla::CreateDenseArray<int32_t>(
119-
{1, 6, std::nullopt, std::nullopt})),
120-
internal::DataItem(schema::kInt32)));
91+
auto fn = test::DataSlice<internal::DataItem>(
92+
{fn_a.item(), fn_b.item(), internal::DataItem(), fn_a.item()},
93+
merged_bag);
94+
auto a_input = test::DataSlice<int>({1, 2, 3, 4});
95+
auto b_input = test::DataSlice<int>({5, 6, 7, std::nullopt});
96+
auto expected1 = test::DataSlice<int>({1, 6, std::nullopt, std::nullopt});
12197
EXPECT_THAT(MapFunctorWithCompilationCache(
12298
fn, /*args=*/{a_input, b_input}, /*kwnames=*/{"a", "b"},
12399
/*include_missing=*/false, /*eval_options=*/{}),
@@ -131,12 +107,7 @@ TEST(MapTest, Basic) {
131107
/*include_missing=*/false, /*eval_options=*/{}),
132108
IsOkAndHolds(IsEquivalentTo(expected1)));
133109

134-
ASSERT_OK_AND_ASSIGN(
135-
DataSlice expected2,
136-
DataSlice::CreateWithFlatShape(
137-
internal::DataSliceImpl::Create(
138-
arolla::CreateDenseArray<int32_t>({1, 6, std::nullopt, 4})),
139-
internal::DataItem(schema::kInt32)));
110+
auto expected2 = test::DataSlice<int>({1, 6, std::nullopt, 4});
140111
EXPECT_THAT(MapFunctorWithCompilationCache(
141112
fn, /*args=*/{a_input, b_input}, /*kwnames=*/{"b"},
142113
/*include_missing=*/true, /*eval_options=*/{}),
@@ -162,11 +133,8 @@ TEST(MapTest, Alignment) {
162133
ASSERT_OK_AND_ASSIGN(auto fn_b,
163134
CreateFunctor(returns_b_expr, koda_signature, {}));
164135
auto merged_bag = DataBag::CommonDataBag({fn_a.GetBag(), fn_b.GetBag()});
165-
ASSERT_OK_AND_ASSIGN(
166-
DataSlice fn,
167-
DataSlice::CreateWithFlatShape(
168-
internal::DataSliceImpl::Create({fn_a.item(), fn_b.item()}),
169-
fn_a.GetSchemaImpl(), merged_bag));
136+
auto fn = test::DataSlice<internal::DataItem>({fn_a.item(), fn_b.item()},
137+
merged_bag);
170138
ASSERT_OK_AND_ASSIGN(auto edge1,
171139
arolla::DenseArrayEdge::FromSplitPoints(
172140
arolla::CreateFullDenseArray<int64_t>({0, 2})));
@@ -175,22 +143,9 @@ TEST(MapTest, Alignment) {
175143
arolla::CreateFullDenseArray<int64_t>({0, 3, 7})));
176144
ASSERT_OK_AND_ASSIGN(auto shape,
177145
DataSlice::JaggedShape::FromEdges({edge1, edge2}));
178-
ASSERT_OK_AND_ASSIGN(
179-
DataSlice a_input,
180-
DataSlice::Create(
181-
internal::DataSliceImpl::Create(
182-
arolla::CreateFullDenseArray<int32_t>({1, 2, 3, 4, 5, 6, 7})),
183-
shape, internal::DataItem(schema::kInt32)));
184-
ASSERT_OK_AND_ASSIGN(DataSlice b_input,
185-
DataSlice::Create(internal::DataItem(8),
186-
internal::DataItem(schema::kInt32)));
187-
188-
ASSERT_OK_AND_ASSIGN(
189-
DataSlice expected,
190-
DataSlice::Create(
191-
internal::DataSliceImpl::Create(
192-
arolla::CreateDenseArray<int32_t>({1, 2, 3, 8, 8, 8, 8})),
193-
shape, internal::DataItem(schema::kInt32)));
146+
auto a_input = test::DataSlice<int>({1, 2, 3, 4, 5, 6, 7}, shape);
147+
auto b_input = test::DataItem(8);
148+
auto expected = test::DataSlice<int>({1, 2, 3, 8, 8, 8, 8}, shape);
194149
EXPECT_THAT(MapFunctorWithCompilationCache(
195150
fn, /*args=*/{a_input, b_input}, /*kwnames=*/{"b"},
196151
/*include_missing=*/false, /*eval_options=*/{}),

‎koladata/functor/signature_storage_test.cc‎

Lines changed: 6 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
#include "koladata/data_bag.h"
2424
#include "koladata/data_slice.h"
2525
#include "koladata/functor/signature.h"
26-
#include "koladata/internal/data_item.h"
2726
#include "koladata/internal/dtype.h"
2827
#include "koladata/object_factories.h"
2928
#include "koladata/operators/masking.h"
@@ -158,9 +157,7 @@ TEST(KodaSignatureToCppSignatureTest, Missing) {
158157
ASSERT_OK_AND_ASSIGN(auto signature, Signature::Create({}));
159158
ASSERT_OK_AND_ASSIGN(auto koda_signature,
160159
CppSignatureToKodaSignature(signature));
161-
ASSERT_OK_AND_ASSIGN(auto missing,
162-
DataSlice::Create(internal::DataItem(arolla::kMissing),
163-
internal::DataItem(schema::kMask)));
160+
auto missing = test::DataItem(arolla::kMissing, schema::kMask);
164161
ASSERT_OK_AND_ASSIGN(auto koda_signature_missing,
165162
ops::ApplyMask(koda_signature, missing));
166163
EXPECT_THAT(
@@ -173,9 +170,7 @@ TEST(KodaSignatureToCppSignatureTest, MissingParameterList) {
173170
ASSERT_OK_AND_ASSIGN(auto koda_signature,
174171
CppSignatureToKodaSignature(signature));
175172
ASSERT_OK_AND_ASSIGN(koda_signature, koda_signature.ForkBag());
176-
ASSERT_OK_AND_ASSIGN(auto missing,
177-
DataSlice::Create(internal::DataItem(arolla::kMissing),
178-
internal::DataItem(schema::kMask)));
173+
auto missing = test::DataItem(arolla::kMissing, schema::kMask);
179174
ASSERT_OK_AND_ASSIGN(auto parameters, koda_signature.GetAttr("parameters"));
180175
ASSERT_OK_AND_ASSIGN(auto parameters_missing,
181176
ops::ApplyMask(parameters, missing));
@@ -194,9 +189,7 @@ TEST(KodaSignatureToCppSignatureTest, MissingParameter) {
194189
ASSERT_OK_AND_ASSIGN(auto koda_signature,
195190
CppSignatureToKodaSignature(signature));
196191
ASSERT_OK_AND_ASSIGN(koda_signature, koda_signature.ForkBag());
197-
ASSERT_OK_AND_ASSIGN(auto missing,
198-
DataSlice::Create(internal::DataItem(arolla::kMissing),
199-
internal::DataItem(schema::kMask)));
192+
auto missing = test::DataItem(arolla::kMissing, schema::kMask);
200193
ASSERT_OK_AND_ASSIGN(auto parameters, koda_signature.GetAttr("parameters"));
201194
ASSERT_OK(parameters.SetInList(test::DataItem(0), missing));
202195
EXPECT_THAT(
@@ -213,9 +206,7 @@ TEST(KodaSignatureToCppSignatureTest, NonTextParameterName) {
213206
ASSERT_OK_AND_ASSIGN(auto koda_signature,
214207
CppSignatureToKodaSignature(signature));
215208
ASSERT_OK_AND_ASSIGN(koda_signature, koda_signature.ForkBag());
216-
ASSERT_OK_AND_ASSIGN(auto missing,
217-
DataSlice::Create(internal::DataItem(arolla::kMissing),
218-
internal::DataItem(schema::kMask)));
209+
auto missing = test::DataItem(arolla::kMissing, schema::kMask);
219210
ASSERT_OK_AND_ASSIGN(auto parameters, koda_signature.GetAttr("parameters"));
220211
ASSERT_OK_AND_ASSIGN(auto parameter,
221212
parameters.GetFromList(test::DataItem(0)));
@@ -238,9 +229,7 @@ TEST(KodaSignatureToCppSignatureTest, DuplicateParameterName) {
238229
ASSERT_OK_AND_ASSIGN(auto koda_signature,
239230
CppSignatureToKodaSignature(signature));
240231
ASSERT_OK_AND_ASSIGN(koda_signature, koda_signature.ForkBag());
241-
ASSERT_OK_AND_ASSIGN(auto missing,
242-
DataSlice::Create(internal::DataItem(arolla::kMissing),
243-
internal::DataItem(schema::kMask)));
232+
auto missing = test::DataItem(arolla::kMissing, schema::kMask);
244233
ASSERT_OK_AND_ASSIGN(auto parameters, koda_signature.GetAttr("parameters"));
245234
ASSERT_OK_AND_ASSIGN(auto parameter,
246235
parameters.GetFromList(test::DataItem(0)));
@@ -259,9 +248,7 @@ TEST(KodaSignatureToCppSignatureTest, InvalidKind) {
259248
ASSERT_OK_AND_ASSIGN(auto koda_signature,
260249
CppSignatureToKodaSignature(signature));
261250
ASSERT_OK_AND_ASSIGN(koda_signature, koda_signature.ForkBag());
262-
ASSERT_OK_AND_ASSIGN(auto missing,
263-
DataSlice::Create(internal::DataItem(arolla::kMissing),
264-
internal::DataItem(schema::kMask)));
251+
auto missing = test::DataItem(arolla::kMissing, schema::kMask);
265252
ASSERT_OK_AND_ASSIGN(auto parameters, koda_signature.GetAttr("parameters"));
266253
ASSERT_OK_AND_ASSIGN(auto parameter,
267254
parameters.GetFromList(test::DataItem(0)));

0 commit comments

Comments
 (0)