Skip to content

Commit 5dd8f16

Browse files
sdf-jklsdf-jkl
committed
move out common tests
1 parent ab9dd3c commit 5dd8f16

4 files changed

Lines changed: 87 additions & 111 deletions

File tree

‎src/oracle/mod.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ pub mod oracle_impl_no_crash;
66
pub mod oracle_impl_tlp_having;
77
pub mod oracle_impl_tlp_where;
88
pub mod oracle_trait;
9+
#[cfg(test)]
10+
pub(crate) mod test_helpers;
911

1012
use std::sync::Arc;
1113

‎src/oracle/oracle_impl_tlp_having.rs‎

Lines changed: 11 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -101,43 +101,15 @@ impl Oracle for TlpHavingOracle {
101101
mod tests {
102102
use super::*;
103103
use crate::common::{FuzzerDataType, LogicalColumn, LogicalTable, init_available_data_types};
104-
use datafusion::arrow::array::{Array, Int64Array};
105-
use datafusion::arrow::datatypes::{DataType, Field, Schema};
106-
use datafusion::arrow::record_batch::RecordBatch;
107-
use datafusion::prelude::SessionContext;
108-
109-
fn make_query_context(label: &str) -> Arc<QueryContext> {
110-
Arc::new(QueryContext::new(
111-
format!("SELECT {}", label),
112-
Arc::new(SessionContext::new()),
113-
))
114-
}
115-
116-
fn make_success_result(label: &str, values: Vec<i64>) -> QueryExecutionResult {
117-
let schema = Arc::new(Schema::new(vec![Field::new("g", DataType::Int64, false)]));
118-
let array = Arc::new(Int64Array::from(values)) as Arc<dyn Array>;
119-
let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
120-
121-
QueryExecutionResult {
122-
query_context: make_query_context(label),
123-
result: Ok(vec![batch]),
124-
}
125-
}
126-
127-
fn make_error_result(label: &str) -> QueryExecutionResult {
128-
QueryExecutionResult {
129-
query_context: make_query_context(label),
130-
result: Err(fuzzer_err("expected execution error in test")),
131-
}
132-
}
104+
use crate::oracle::test_helpers;
133105

134106
#[tokio::test]
135107
async fn tlp_having_validate_passes_for_matching_values() {
136108
let oracle =
137109
TlpHavingOracle::new(1, Arc::new(crate::fuzz_context::GlobalContext::default()));
138110
let results = vec![
139-
make_success_result("all", vec![1, 2, 3]),
140-
make_success_result("partition_union", vec![1, 2, 3]),
111+
test_helpers::make_success_result("all", "g", vec![1, 2, 3]),
112+
test_helpers::make_success_result("partition_union", "g", vec![1, 2, 3]),
141113
];
142114

143115
assert!(oracle.validate_consistency(&results).await.is_ok());
@@ -148,8 +120,8 @@ mod tests {
148120
let oracle =
149121
TlpHavingOracle::new(1, Arc::new(crate::fuzz_context::GlobalContext::default()));
150122
let results = vec![
151-
make_success_result("all", vec![1, 2]),
152-
make_success_result("partition_union", vec![1, 2, 2]),
123+
test_helpers::make_success_result("all", "g", vec![1, 2]),
124+
test_helpers::make_success_result("partition_union", "g", vec![1, 2, 2]),
153125
];
154126

155127
let err = oracle.validate_consistency(&results).await.unwrap_err();
@@ -164,44 +136,27 @@ mod tests {
164136
let oracle =
165137
TlpHavingOracle::new(1, Arc::new(crate::fuzz_context::GlobalContext::default()));
166138
let results = vec![
167-
make_success_result("all", vec![1, 2]),
168-
make_error_result("partition_union"),
139+
test_helpers::make_success_result("all", "g", vec![1, 2]),
140+
test_helpers::make_error_result("partition_union"),
169141
];
170142

171143
assert!(oracle.validate_consistency(&results).await.is_ok());
172144
}
173145

174146
#[tokio::test]
175147
async fn tlp_having_validate_fails_for_schema_mismatch() {
176-
let one_col_schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int64, false)]));
177-
let one_col_batch = RecordBatch::try_new(
178-
one_col_schema,
179-
vec![Arc::new(Int64Array::from(vec![1, 2])) as Arc<dyn Array>],
180-
)
181-
.unwrap();
182-
183-
let two_col_schema = Arc::new(Schema::new(vec![
184-
Field::new("c1", DataType::Int64, false),
185-
Field::new("c2", DataType::Int64, false),
186-
]));
187-
let two_col_batch = RecordBatch::try_new(
188-
two_col_schema,
189-
vec![
190-
Arc::new(Int64Array::from(vec![1])) as Arc<dyn Array>,
191-
Arc::new(Int64Array::from(vec![9])) as Arc<dyn Array>,
192-
],
193-
)
194-
.unwrap();
148+
let one_col_batch = test_helpers::make_one_col_batch(vec![1, 2]);
149+
let two_col_batch = test_helpers::make_two_col_batch(vec![1], vec![9]);
195150

196151
let oracle =
197152
TlpHavingOracle::new(1, Arc::new(crate::fuzz_context::GlobalContext::default()));
198153
let results = vec![
199154
QueryExecutionResult {
200-
query_context: make_query_context("all"),
155+
query_context: test_helpers::make_query_context("all"),
201156
result: Ok(vec![one_col_batch.clone()]),
202157
},
203158
QueryExecutionResult {
204-
query_context: make_query_context("partition_union"),
159+
query_context: test_helpers::make_query_context("partition_union"),
205160
result: Ok(vec![two_col_batch]),
206161
},
207162
];

‎src/oracle/oracle_impl_tlp_where.rs‎

Lines changed: 11 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -100,42 +100,15 @@ impl Oracle for TlpWhereOracle {
100100
mod tests {
101101
use super::*;
102102
use crate::common::{LogicalTable, init_available_data_types};
103-
use datafusion::arrow::array::{Array, Int64Array, RecordBatch};
104-
use datafusion::arrow::datatypes::{DataType, Field, Schema};
105-
use datafusion::prelude::SessionContext;
106-
107-
fn make_query_context(label: &str) -> Arc<QueryContext> {
108-
Arc::new(QueryContext::new(
109-
format!("SELECT {}", label),
110-
Arc::new(SessionContext::new()),
111-
))
112-
}
113-
114-
fn make_success_result(label: &str, values: Vec<i64>) -> QueryExecutionResult {
115-
let schema = Arc::new(Schema::new(vec![Field::new("cnt", DataType::Int64, false)]));
116-
let array = Arc::new(Int64Array::from(values)) as Arc<dyn Array>;
117-
let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
118-
119-
QueryExecutionResult {
120-
query_context: make_query_context(label),
121-
result: Ok(vec![batch]),
122-
}
123-
}
124-
125-
fn make_error_result(label: &str) -> QueryExecutionResult {
126-
QueryExecutionResult {
127-
query_context: make_query_context(label),
128-
result: Err(fuzzer_err("expected execution error in test")),
129-
}
130-
}
103+
use crate::oracle::test_helpers;
131104

132105
#[tokio::test]
133106
async fn tlp_where_validate_passes_for_matching_values() {
134107
let oracle =
135108
TlpWhereOracle::new(1, Arc::new(crate::fuzz_context::GlobalContext::default()));
136109
let results = vec![
137-
make_success_result("all", vec![1, 2, 3]),
138-
make_success_result("partition_union", vec![1, 2, 3]),
110+
test_helpers::make_success_result("all", "cnt", vec![1, 2, 3]),
111+
test_helpers::make_success_result("partition_union", "cnt", vec![1, 2, 3]),
139112
];
140113

141114
assert!(oracle.validate_consistency(&results).await.is_ok());
@@ -146,8 +119,8 @@ mod tests {
146119
let oracle =
147120
TlpWhereOracle::new(1, Arc::new(crate::fuzz_context::GlobalContext::default()));
148121
let results = vec![
149-
make_success_result("all", vec![1, 2]),
150-
make_success_result("partition_union", vec![1, 2, 2]),
122+
test_helpers::make_success_result("all", "cnt", vec![1, 2]),
123+
test_helpers::make_success_result("partition_union", "cnt", vec![1, 2, 2]),
151124
];
152125

153126
let err = oracle.validate_consistency(&results).await.unwrap_err();
@@ -162,44 +135,27 @@ mod tests {
162135
let oracle =
163136
TlpWhereOracle::new(1, Arc::new(crate::fuzz_context::GlobalContext::default()));
164137
let results = vec![
165-
make_success_result("all", vec![1, 2]),
166-
make_error_result("partition_union"),
138+
test_helpers::make_success_result("all", "cnt", vec![1, 2]),
139+
test_helpers::make_error_result("partition_union"),
167140
];
168141

169142
assert!(oracle.validate_consistency(&results).await.is_ok());
170143
}
171144

172145
#[tokio::test]
173146
async fn tlp_where_validate_fails_for_schema_mismatch() {
174-
let one_col_schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int64, false)]));
175-
let one_col_batch = RecordBatch::try_new(
176-
one_col_schema,
177-
vec![Arc::new(Int64Array::from(vec![1, 2])) as Arc<dyn Array>],
178-
)
179-
.unwrap();
180-
181-
let two_col_schema = Arc::new(Schema::new(vec![
182-
Field::new("c1", DataType::Int64, false),
183-
Field::new("c2", DataType::Int64, false),
184-
]));
185-
let two_col_batch = RecordBatch::try_new(
186-
two_col_schema,
187-
vec![
188-
Arc::new(Int64Array::from(vec![1])) as Arc<dyn Array>,
189-
Arc::new(Int64Array::from(vec![9])) as Arc<dyn Array>,
190-
],
191-
)
192-
.unwrap();
147+
let one_col_batch = test_helpers::make_one_col_batch(vec![1, 2]);
148+
let two_col_batch = test_helpers::make_two_col_batch(vec![1], vec![9]);
193149

194150
let oracle =
195151
TlpWhereOracle::new(1, Arc::new(crate::fuzz_context::GlobalContext::default()));
196152
let results = vec![
197153
QueryExecutionResult {
198-
query_context: make_query_context("all"),
154+
query_context: test_helpers::make_query_context("all"),
199155
result: Ok(vec![one_col_batch.clone()]),
200156
},
201157
QueryExecutionResult {
202-
query_context: make_query_context("partition_union"),
158+
query_context: test_helpers::make_query_context("partition_union"),
203159
result: Ok(vec![two_col_batch]),
204160
},
205161
];

‎src/oracle/test_helpers.rs‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
use crate::common::fuzzer_err;
2+
use crate::oracle::{QueryContext, QueryExecutionResult};
3+
use datafusion::arrow::array::{Array, Int64Array, RecordBatch};
4+
use datafusion::arrow::datatypes::{DataType, Field, Schema};
5+
use datafusion::prelude::SessionContext;
6+
use std::sync::Arc;
7+
8+
pub(crate) fn make_query_context(label: &str) -> Arc<QueryContext> {
9+
Arc::new(QueryContext::new(
10+
format!("SELECT {}", label),
11+
Arc::new(SessionContext::new()),
12+
))
13+
}
14+
15+
pub(crate) fn make_success_result(
16+
label: &str,
17+
column_name: &str,
18+
values: Vec<i64>,
19+
) -> QueryExecutionResult {
20+
let schema = Arc::new(Schema::new(vec![Field::new(
21+
column_name,
22+
DataType::Int64,
23+
false,
24+
)]));
25+
let array = Arc::new(Int64Array::from(values)) as Arc<dyn Array>;
26+
let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
27+
28+
QueryExecutionResult {
29+
query_context: make_query_context(label),
30+
result: Ok(vec![batch]),
31+
}
32+
}
33+
34+
pub(crate) fn make_error_result(label: &str) -> QueryExecutionResult {
35+
QueryExecutionResult {
36+
query_context: make_query_context(label),
37+
result: Err(fuzzer_err("expected execution error in test")),
38+
}
39+
}
40+
41+
pub(crate) fn make_one_col_batch(values: Vec<i64>) -> RecordBatch {
42+
let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int64, false)]));
43+
RecordBatch::try_new(
44+
schema,
45+
vec![Arc::new(Int64Array::from(values)) as Arc<dyn Array>],
46+
)
47+
.unwrap()
48+
}
49+
50+
pub(crate) fn make_two_col_batch(col1: Vec<i64>, col2: Vec<i64>) -> RecordBatch {
51+
let schema = Arc::new(Schema::new(vec![
52+
Field::new("c1", DataType::Int64, false),
53+
Field::new("c2", DataType::Int64, false),
54+
]));
55+
RecordBatch::try_new(
56+
schema,
57+
vec![
58+
Arc::new(Int64Array::from(col1)) as Arc<dyn Array>,
59+
Arc::new(Int64Array::from(col2)) as Arc<dyn Array>,
60+
],
61+
)
62+
.unwrap()
63+
}

0 commit comments

Comments
 (0)