Skip to content

Commit

Permalink
test: cleanup test definitions with rstest
Browse files Browse the repository at this point in the history
  • Loading branch information
@aljazerzen
aljazerzen committed Feb 20, 2024
1 parent 7d6a05c commit 1e3d842
Show file tree
Hide file tree
Showing 10 changed files with 327 additions and 352 deletions.
55 changes: 55 additions & 0 deletions Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions connector_arrow/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ similar-asserts = { version = "1.5.0" }
half = "2.3.1"
rand = { version = "0.8.5", default-features = false }
rand_chacha = "0.3.1"
rstest = { version = "0.18.2", default-features = false }


[features]
Expand Down
207 changes: 3 additions & 204 deletions connector_arrow/tests/it/generator.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,10 @@ use half::f16;
use rand::Rng;
use std::sync::Arc;

use super::spec::*;

pub fn generate_batch<R: Rng>(
column_specs: Vec<ColumnSpec>,
column_specs: ArrowGenSpec,
rng: &mut R,
) -> (SchemaRef, Vec<RecordBatch>) {
let mut arrays = Vec::new();
Expand All @@ -27,209 +29,6 @@ pub fn generate_batch<R: Rng>(
}
}

pub fn spec_all_types() -> Vec<ColumnSpec> {
domains_to_batch_spec(
&[
DataType::Null,
DataType::Boolean,
// DataType::Int8,
DataType::Int16,
DataType::Int32,
DataType::Int64,
// DataType::UInt8,
// DataType::UInt16,
// DataType::UInt32,
// DataType::UInt64,
// DataType::Float16,
DataType::Float32,
DataType::Float64,
// DataType::Timestamp(TimeUnit::Nanosecond, None),
// DataType::Timestamp(TimeUnit::Microsecond, None),
// DataType::Timestamp(TimeUnit::Millisecond, None),
// DataType::Timestamp(TimeUnit::Second, None),
// DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("+07:30"))),
// DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("+07:30"))),
// DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from("+07:30"))),
// DataType::Timestamp(TimeUnit::Second, Some(Arc::from("+07:30"))),
// DataType::Time32(TimeUnit::Millisecond),
// DataType::Time32(TimeUnit::Second),
// DataType::Time64(TimeUnit::Nanosecond),
// DataType::Time64(TimeUnit::Microsecond),
// DataType::Duration(TimeUnit::Nanosecond),
// DataType::Duration(TimeUnit::Microsecond),
// DataType::Duration(TimeUnit::Millisecond),
// DataType::Duration(TimeUnit::Second),
// DataType::Interval(IntervalUnit::YearMonth),
// DataType::Interval(IntervalUnit::MonthDayNano),
// DataType::Interval(IntervalUnit::DayTime),
],
&[false, true],
&[ValueGenProcess::High],
)
}

pub fn spec_empty() -> Vec<ColumnSpec> {
domains_to_batch_spec(
&[DataType::Null, DataType::Int64, DataType::Float64],
&[false, true],
&[],
)
}

pub fn spec_null_bool() -> Vec<ColumnSpec> {
domains_to_batch_spec(
&[DataType::Null, DataType::Boolean],
&[false, true],
&VALUE_GEN_PROCESS_ALL,
)
}

pub fn spec_numeric() -> Vec<ColumnSpec> {
domains_to_batch_spec(
&[
DataType::Int8,
DataType::Int16,
DataType::Int32,
DataType::Int64,
DataType::UInt8,
DataType::UInt16,
DataType::UInt32,
DataType::UInt64,
DataType::Float16,
DataType::Float32,
DataType::Float64,
],
&[false, true],
&VALUE_GEN_PROCESS_ALL,
)
}

pub fn spec_timestamp() -> Vec<ColumnSpec> {
domains_to_batch_spec(
&[
DataType::Timestamp(TimeUnit::Nanosecond, None),
DataType::Timestamp(TimeUnit::Microsecond, None),
DataType::Timestamp(TimeUnit::Millisecond, None),
DataType::Timestamp(TimeUnit::Second, None),
DataType::Timestamp(TimeUnit::Nanosecond, Some(Arc::from("+07:30"))),
DataType::Timestamp(TimeUnit::Microsecond, Some(Arc::from("+07:30"))),
DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from("+07:30"))),
DataType::Timestamp(TimeUnit::Second, Some(Arc::from("+07:30"))),
],
&[true],
&VALUE_GEN_PROCESS_ALL,
)
}
pub fn spec_date() -> Vec<ColumnSpec> {
domains_to_batch_spec(
&[DataType::Date32, DataType::Date64],
&[true],
&VALUE_GEN_PROCESS_ALL,
)
}
pub fn spec_time() -> Vec<ColumnSpec> {
domains_to_batch_spec(
&[
DataType::Time32(TimeUnit::Millisecond),
DataType::Time32(TimeUnit::Second),
DataType::Time64(TimeUnit::Nanosecond),
DataType::Time64(TimeUnit::Microsecond),
],
&[true],
&VALUE_GEN_PROCESS_ALL,
)
}
pub fn spec_duration() -> Vec<ColumnSpec> {
domains_to_batch_spec(
&[
DataType::Duration(TimeUnit::Nanosecond),
DataType::Duration(TimeUnit::Microsecond),
DataType::Duration(TimeUnit::Millisecond),
DataType::Duration(TimeUnit::Second),
],
&[true],
&VALUE_GEN_PROCESS_ALL,
)
}
pub fn spec_interval() -> Vec<ColumnSpec> {
domains_to_batch_spec(
&[
DataType::Interval(IntervalUnit::YearMonth),
DataType::Interval(IntervalUnit::MonthDayNano),
DataType::Interval(IntervalUnit::DayTime),
],
&[true],
&VALUE_GEN_PROCESS_ALL,
)
}

pub fn domains_to_batch_spec(
data_types_domain: &[DataType],
is_nullable_domain: &[bool],
value_gen_process_domain: &[ValueGenProcess],
) -> Vec<ColumnSpec> {
let mut columns = Vec::new();
for data_type in data_types_domain {
for is_nullable in is_nullable_domain {
let is_nullable = *is_nullable;
if matches!(data_type, &DataType::Null) && !is_nullable {
continue;
}

let mut field_name = data_type.to_string();
if is_nullable {
field_name += "_null";
}
let mut col = ColumnSpec {
field_name,
data_type: data_type.clone(),
is_nullable,
values: Vec::new(),
};

for gen_process in value_gen_process_domain {
col.values.push(ValuesSpec {
gen_process: if matches!(gen_process, ValueGenProcess::Null) && !is_nullable {
ValueGenProcess::RandomUniform
} else {
*gen_process
},
repeat: 1,
});
}
columns.push(col);
}
}
columns
}

#[derive(Clone, Copy)]
pub enum ValueGenProcess {
Null,
Low,
High,
RandomUniform,
}

const VALUE_GEN_PROCESS_ALL: [ValueGenProcess; 4] = [
ValueGenProcess::Low,
ValueGenProcess::High,
ValueGenProcess::Null,
ValueGenProcess::RandomUniform,
];

struct ValuesSpec {
gen_process: ValueGenProcess,
repeat: usize,
}

pub struct ColumnSpec {
field_name: String,
is_nullable: bool,
data_type: DataType,
values: Vec<ValuesSpec>,
}

fn count_values(values: &[ValuesSpec]) -> usize {
values.iter().map(|v| v.repeat).sum()
}
Expand Down
1 change: 1 addition & 0 deletions connector_arrow/tests/it/main.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
mod generator;
mod spec;
mod tests;
mod util;

Expand Down
Loading

0 comments on commit 1e3d842

Please sign in to comment.