Skip to content

Commit c75705e

Browse files
Sl1mb0Sl1mb0
committed
fix: broken tests
1 parent 854ce39 commit c75705e

File tree

5 files changed

+103
-78
lines changed

5 files changed

+103
-78
lines changed

‎host/tests/integration_tests/mod.rs‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,2 @@
11
mod python;
22
mod rust;
3-
mod test_utils;

‎host/tests/integration_tests/python/runtime/http.rs‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,7 @@ use tokio::runtime::Handle;
1717
use wasmtime_wasi_http::types::DEFAULT_FORBIDDEN_HEADERS;
1818
use wiremock::{Mock, MockServer, ResponseTemplate, matchers};
1919

20-
use crate::integration_tests::{
21-
python::test_utils::{python_component, python_scalar_udf},
22-
test_utils::ColumnarValueExt,
23-
};
20+
use crate::integration_tests::python::test_utils::{python_component, python_scalar_udf};
2421

2522
#[tokio::test(flavor = "multi_thread")]
2623
async fn test_requests_simple() {
@@ -658,14 +655,17 @@ def perform_request(url: str) -> str:
658655
rt_tmp.shutdown_timeout(Duration::from_secs(1));
659656

660657
let array = rt_cpu.block_on(async {
661-
udf.invoke_with_args(ScalarFunctionArgs {
662-
args: vec![ColumnarValue::Scalar(ScalarValue::Utf8(Some(server.uri())))],
663-
arg_fields: vec![Arc::new(Field::new("uri", DataType::Utf8, true))],
664-
number_rows: 1,
665-
return_field: Arc::new(Field::new("r", DataType::Utf8, true)),
666-
})
658+
udf.invoke_async_with_args(
659+
ScalarFunctionArgs {
660+
args: vec![ColumnarValue::Scalar(ScalarValue::Utf8(Some(server.uri())))],
661+
arg_fields: vec![Arc::new(Field::new("uri", DataType::Utf8, true))],
662+
number_rows: 1,
663+
return_field: Arc::new(Field::new("r", DataType::Utf8, true)),
664+
},
665+
&ConfigOptions::default(),
666+
)
667+
.await
667668
.unwrap()
668-
.unwrap_array()
669669
});
670670

671671
assert_eq!(

‎host/tests/integration_tests/python/types/none.rs‎

Lines changed: 80 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,14 @@ use arrow::{
44
array::{Array, NullArray},
55
datatypes::{DataType, Field},
66
};
7-
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
8-
9-
use crate::integration_tests::{
10-
python::test_utils::python_scalar_udf, test_utils::ColumnarValueExt,
7+
use datafusion_common::config::ConfigOptions;
8+
use datafusion_expr::{
9+
ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
10+
async_udf::AsyncScalarUDFImpl,
1111
};
1212

13+
use crate::integration_tests::python::test_utils::python_scalar_udf;
14+
1315
#[tokio::test(flavor = "multi_thread")]
1416
async fn test_ok() {
1517
const CODE: &str = "
@@ -26,14 +28,17 @@ def foo(x: None) -> None:
2628
assert_eq!(udf.return_type(&[DataType::Null]).unwrap(), DataType::Null,);
2729

2830
let array = udf
29-
.invoke_with_args(ScalarFunctionArgs {
30-
args: vec![ColumnarValue::Array(Arc::new(NullArray::new(3)))],
31-
arg_fields: vec![Arc::new(Field::new("a1", DataType::Null, true))],
32-
number_rows: 3,
33-
return_field: Arc::new(Field::new("r", DataType::Null, true)),
34-
})
35-
.unwrap()
36-
.unwrap_array();
31+
.invoke_async_with_args(
32+
ScalarFunctionArgs {
33+
args: vec![ColumnarValue::Array(Arc::new(NullArray::new(3)))],
34+
arg_fields: vec![Arc::new(Field::new("a1", DataType::Null, true))],
35+
number_rows: 3,
36+
return_field: Arc::new(Field::new("r", DataType::Null, true)),
37+
},
38+
&ConfigOptions::default(),
39+
)
40+
.await
41+
.unwrap();
3742
assert_eq!(array.as_ref(), &NullArray::new(3) as &dyn Array,);
3843
}
3944

@@ -46,14 +51,17 @@ def foo() -> None:
4651
let udf = python_scalar_udf(CODE).await.unwrap();
4752

4853
let array = udf
49-
.invoke_with_args(ScalarFunctionArgs {
50-
args: vec![],
51-
arg_fields: vec![],
52-
number_rows: 3,
53-
return_field: Arc::new(Field::new("r", DataType::Null, true)),
54-
})
55-
.unwrap()
56-
.unwrap_array();
54+
.invoke_async_with_args(
55+
ScalarFunctionArgs {
56+
args: vec![],
57+
arg_fields: vec![],
58+
number_rows: 3,
59+
return_field: Arc::new(Field::new("r", DataType::Null, true)),
60+
},
61+
&ConfigOptions::default(),
62+
)
63+
.await
64+
.unwrap();
5765
assert_eq!(array.as_ref(), &NullArray::new(3) as &dyn Array,);
5866
}
5967

@@ -66,14 +74,17 @@ def foo() -> None:
6674
let udf = python_scalar_udf(CODE).await.unwrap();
6775

6876
let array = udf
69-
.invoke_with_args(ScalarFunctionArgs {
70-
args: vec![],
71-
arg_fields: vec![],
72-
number_rows: 3,
73-
return_field: Arc::new(Field::new("r", DataType::Null, true)),
74-
})
75-
.unwrap()
76-
.unwrap_array();
77+
.invoke_async_with_args(
78+
ScalarFunctionArgs {
79+
args: vec![],
80+
arg_fields: vec![],
81+
number_rows: 3,
82+
return_field: Arc::new(Field::new("r", DataType::Null, true)),
83+
},
84+
&ConfigOptions::default(),
85+
)
86+
.await
87+
.unwrap();
7788
assert_eq!(array.as_ref(), &NullArray::new(3) as &dyn Array,);
7889
}
7990

@@ -88,14 +99,17 @@ def foo() -> None:
8899
let udf = python_scalar_udf(CODE).await.unwrap();
89100

90101
let array = udf
91-
.invoke_with_args(ScalarFunctionArgs {
92-
args: vec![],
93-
arg_fields: vec![],
94-
number_rows: 3,
95-
return_field: Arc::new(Field::new("r", DataType::Null, true)),
96-
})
97-
.unwrap()
98-
.unwrap_array();
102+
.invoke_async_with_args(
103+
ScalarFunctionArgs {
104+
args: vec![],
105+
arg_fields: vec![],
106+
number_rows: 3,
107+
return_field: Arc::new(Field::new("r", DataType::Null, true)),
108+
},
109+
&ConfigOptions::default(),
110+
)
111+
.await
112+
.unwrap();
99113
assert_eq!(array.as_ref(), &NullArray::new(3) as &dyn Array,);
100114
}
101115

@@ -108,12 +122,16 @@ def foo() -> None:
108122
let udf = python_scalar_udf(CODE).await.unwrap();
109123

110124
let err = udf
111-
.invoke_with_args(ScalarFunctionArgs {
112-
args: vec![],
113-
arg_fields: vec![],
114-
number_rows: 3,
115-
return_field: Arc::new(Field::new("r", DataType::Null, true)),
116-
})
125+
.invoke_async_with_args(
126+
ScalarFunctionArgs {
127+
args: vec![],
128+
arg_fields: vec![],
129+
number_rows: 3,
130+
return_field: Arc::new(Field::new("r", DataType::Null, true)),
131+
},
132+
&ConfigOptions::default(),
133+
)
134+
.await
117135
.unwrap_err();
118136

119137
insta::assert_snapshot!(
@@ -131,12 +149,16 @@ def foo() -> None:
131149
let udf = python_scalar_udf(CODE).await.unwrap();
132150

133151
let err = udf
134-
.invoke_with_args(ScalarFunctionArgs {
135-
args: vec![],
136-
arg_fields: vec![],
137-
number_rows: 3,
138-
return_field: Arc::new(Field::new("r", DataType::Null, true)),
139-
})
152+
.invoke_async_with_args(
153+
ScalarFunctionArgs {
154+
args: vec![],
155+
arg_fields: vec![],
156+
number_rows: 3,
157+
return_field: Arc::new(Field::new("r", DataType::Null, true)),
158+
},
159+
&ConfigOptions::default(),
160+
)
161+
.await
140162
.unwrap_err();
141163

142164
insta::assert_snapshot!(
@@ -154,12 +176,16 @@ def foo() -> None:
154176
let udf = python_scalar_udf(CODE).await.unwrap();
155177

156178
let err = udf
157-
.invoke_with_args(ScalarFunctionArgs {
158-
args: vec![],
159-
arg_fields: vec![],
160-
number_rows: 3,
161-
return_field: Arc::new(Field::new("r", DataType::Null, true)),
162-
})
179+
.invoke_async_with_args(
180+
ScalarFunctionArgs {
181+
args: vec![],
182+
arg_fields: vec![],
183+
number_rows: 3,
184+
return_field: Arc::new(Field::new("r", DataType::Null, true)),
185+
},
186+
&ConfigOptions::default(),
187+
)
188+
.await
163189
.unwrap_err();
164190

165191
insta::assert_snapshot!(

‎host/tests/integration_tests/rust.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,6 @@ async fn test_add_one() {
7979
.unwrap();
8080
assert_eq!(
8181
array.as_ref(),
82-
&Int32Array::from_iter([Some(4)]) as &dyn Array,
82+
&Int32Array::from_iter([Some(4), Some(4), Some(4)]) as &dyn Array,
8383
);
8484
}

‎query/tests/integration.rs‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -274,17 +274,17 @@ EXPLAIN SELECT add_one(1);
274274
insta::assert_snapshot!(
275275
batches_to_string(&batch),
276276
@r"
277-
│++---------------+------------------------------------------------------------------------------+
278-
│+| plan_type | plan |
279-
│++---------------+------------------------------------------------------------------------------+
280-
│+| logical_plan | Projection: add_one(Int64(1)) |
281-
│+| | EmptyRelation |
282-
│+| physical_plan | ProjectionExec: expr=[__async_fn_0@0 as add_one(Int64(1))] |
283-
│+| | AsyncFuncExec: async_expr=[async_expr(name=__async_fn_0, expr=add_one(1))] |
284-
│+| | CoalesceBatchesExec: target_batch_size=8192 |
285-
│+| | PlaceholderRowExec |
286-
│+| | |
287-
│++---------------+------------------------------------------------------------------------------+
277+
+---------------+------------------------------------------------------------------------------+
278+
| plan_type | plan |
279+
+---------------+------------------------------------------------------------------------------+
280+
| logical_plan | Projection: add_one(Int64(1)) |
281+
| | EmptyRelation |
282+
| physical_plan | ProjectionExec: expr=[__async_fn_0@0 as add_one(Int64(1))] |
283+
| | AsyncFuncExec: async_expr=[async_expr(name=__async_fn_0, expr=add_one(1))] |
284+
| | CoalesceBatchesExec: target_batch_size=8192 |
285+
| | PlaceholderRowExec |
286+
| | |
287+
+---------------+------------------------------------------------------------------------------+
288288
");
289289
}
290290

0 commit comments

Comments
 (0)