refactor: Extract a helper fn for test messages

Reduces duplication, makes it easier to see what's different between the
tests, will make it easier to add another field in the next commit
pull/24376/head
Carol (Nichols || Goulding) 2022-12-09 10:39:21 -05:00
parent e91d8998a8
commit 33886970ef
No known key found for this signature in database
GPG Key ID: E907EE5A736F87D4
1 changed files with 63 additions and 89 deletions

View File

@ -1392,16 +1392,12 @@ mod tests {
MockFlightClient::new([( MockFlightClient::new([(
"addr1", "addr1",
Ok(MockQueryData { Ok(MockQueryData {
results: vec![Ok(( results: vec![metadata(
LowLevelMessage::None, 1,
IngesterQueryResponseMetadata { Some(PartitionStatus {
partition_id: 1,
status: Some(PartitionStatus {
parquet_max_sequence_number: None, parquet_max_sequence_number: None,
}), }),
ingester_uuid: String::new(), )],
},
))],
}), }),
)]) )])
.await, .await,
@ -1425,14 +1421,7 @@ mod tests {
MockFlightClient::new([( MockFlightClient::new([(
"addr1", "addr1",
Ok(MockQueryData { Ok(MockQueryData {
results: vec![Ok(( results: vec![metadata(1, None)],
LowLevelMessage::None,
IngesterQueryResponseMetadata {
partition_id: 1,
status: None,
ingester_uuid: String::new(),
},
))],
}), }),
)]) )])
.await, .await,
@ -1449,36 +1438,24 @@ mod tests {
"addr1", "addr1",
Ok(MockQueryData { Ok(MockQueryData {
results: vec![ results: vec![
Ok(( metadata(
LowLevelMessage::None, 1,
IngesterQueryResponseMetadata { Some(PartitionStatus {
partition_id: 1,
status: Some(PartitionStatus {
parquet_max_sequence_number: None, parquet_max_sequence_number: None,
}), }),
ingester_uuid: String::new(), ),
}, metadata(
)), 2,
Ok(( Some(PartitionStatus {
LowLevelMessage::None,
IngesterQueryResponseMetadata {
partition_id: 2,
status: Some(PartitionStatus {
parquet_max_sequence_number: None, parquet_max_sequence_number: None,
}), }),
ingester_uuid: String::new(), ),
}, metadata(
)), 1,
Ok(( Some(PartitionStatus {
LowLevelMessage::None,
IngesterQueryResponseMetadata {
partition_id: 1,
status: Some(PartitionStatus {
parquet_max_sequence_number: None, parquet_max_sequence_number: None,
}), }),
ingester_uuid: String::new(), ),
},
)),
], ],
}), }),
)]) )])
@ -1549,16 +1526,12 @@ mod tests {
"addr1", "addr1",
Ok(MockQueryData { Ok(MockQueryData {
results: vec![ results: vec![
Ok(( metadata(
LowLevelMessage::None, 1,
IngesterQueryResponseMetadata { Some(PartitionStatus {
partition_id: 1,
status: Some(PartitionStatus {
parquet_max_sequence_number: Some(11), parquet_max_sequence_number: Some(11),
}), }),
ingester_uuid: String::new(), ),
},
)),
Ok(( Ok((
LowLevelMessage::Schema(Arc::clone(&schema_1_1)), LowLevelMessage::Schema(Arc::clone(&schema_1_1)),
IngesterQueryResponseMetadata::default(), IngesterQueryResponseMetadata::default(),
@ -1579,16 +1552,12 @@ mod tests {
LowLevelMessage::RecordBatch(record_batch_1_2), LowLevelMessage::RecordBatch(record_batch_1_2),
IngesterQueryResponseMetadata::default(), IngesterQueryResponseMetadata::default(),
)), )),
Ok(( metadata(
LowLevelMessage::None, 2,
IngesterQueryResponseMetadata { Some(PartitionStatus {
partition_id: 2,
status: Some(PartitionStatus {
parquet_max_sequence_number: Some(21), parquet_max_sequence_number: Some(21),
}), }),
ingester_uuid: String::new(), ),
},
)),
Ok(( Ok((
LowLevelMessage::Schema(Arc::clone(&schema_2_1)), LowLevelMessage::Schema(Arc::clone(&schema_2_1)),
IngesterQueryResponseMetadata::default(), IngesterQueryResponseMetadata::default(),
@ -1604,16 +1573,12 @@ mod tests {
"addr2", "addr2",
Ok(MockQueryData { Ok(MockQueryData {
results: vec![ results: vec![
Ok(( metadata(
LowLevelMessage::None, 3,
IngesterQueryResponseMetadata { Some(PartitionStatus {
partition_id: 3,
status: Some(PartitionStatus {
parquet_max_sequence_number: Some(31), parquet_max_sequence_number: Some(31),
}), }),
ingester_uuid: String::new(), ),
},
)),
Ok(( Ok((
LowLevelMessage::Schema(Arc::clone(&schema_3_1)), LowLevelMessage::Schema(Arc::clone(&schema_3_1)),
IngesterQueryResponseMetadata::default(), IngesterQueryResponseMetadata::default(),
@ -1784,16 +1749,12 @@ mod tests {
"addr1", "addr1",
Ok(MockQueryData { Ok(MockQueryData {
results: vec![ results: vec![
Ok(( metadata(
LowLevelMessage::None, 1,
IngesterQueryResponseMetadata { Some(PartitionStatus {
partition_id: 1,
status: Some(PartitionStatus {
parquet_max_sequence_number: Some(11), parquet_max_sequence_number: Some(11),
}), }),
ingester_uuid: String::new(), ),
},
)),
Ok(( Ok((
LowLevelMessage::Schema(Arc::clone(&schema_1_1)), LowLevelMessage::Schema(Arc::clone(&schema_1_1)),
IngesterQueryResponseMetadata::default(), IngesterQueryResponseMetadata::default(),
@ -1876,9 +1837,22 @@ mod tests {
lp_to_mutable_batch(lp).1.to_arrow(Projection::All).unwrap() lp_to_mutable_batch(lp).1.to_arrow(Projection::All).unwrap()
} }
type MockFlightResult = Result<(LowLevelMessage, IngesterQueryResponseMetadata), FlightError>;
fn metadata(partition_id: i64, status: Option<PartitionStatus>) -> MockFlightResult {
Ok((
LowLevelMessage::None,
IngesterQueryResponseMetadata {
partition_id,
status,
ingester_uuid: String::new(),
},
))
}
#[derive(Debug)] #[derive(Debug)]
struct MockQueryData { struct MockQueryData {
results: Vec<Result<(LowLevelMessage, IngesterQueryResponseMetadata), FlightError>>, results: Vec<MockFlightResult>,
} }
#[async_trait] #[async_trait]