From d8c512fe4216834869a5fc60e1b5aa1903d1e3eb Mon Sep 17 00:00:00 2001 From: Jefffrey Date: Wed, 28 Feb 2024 21:39:23 +1100 Subject: [PATCH] Use assert_batches_eq in tests --- tests/basic/main.rs | 200 +++++++++++++++++++++----------------------- 1 file changed, 95 insertions(+), 105 deletions(-) diff --git a/tests/basic/main.rs b/tests/basic/main.rs index 63c7016d..6d9c37c2 100644 --- a/tests/basic/main.rs +++ b/tests/basic/main.rs @@ -139,20 +139,18 @@ pub fn basic_test() { let reader = new_arrow_reader(&path, &["a", "b", "str_direct", "d", "e", "f"]); let batch = reader.collect::, _>>().unwrap(); - let expected = r#"+-----+-------+------------+-----+-----+-------+ -| a | b | str_direct | d | e | f | -+-----+-------+------------+-----+-----+-------+ -| 1.0 | true | a | a | ddd | aaaaa | -| 2.0 | false | cccccc | bb | cc | bbbbb | -| | | | | | | -| 4.0 | true | ddd | ccc | bb | ccccc | -| 5.0 | false | ee | ddd | a | ddddd | -+-----+-------+------------+-----+-----+-------+"#; - - assert_eq!( - expected, - pretty::pretty_format_batches(&batch).unwrap().to_string() - ) + let expected = [ + "+-----+-------+------------+-----+-----+-------+", + "| a | b | str_direct | d | e | f |", + "+-----+-------+------------+-----+-----+-------+", + "| 1.0 | true | a | a | ddd | aaaaa |", + "| 2.0 | false | cccccc | bb | cc | bbbbb |", + "| | | | | | |", + "| 4.0 | true | ddd | ccc | bb | ccccc |", + "| 5.0 | false | ee | ddd | a | ddddd |", + "+-----+-------+------------+-----+-----+-------+", + ]; + assert_batches_eq(&batch, &expected); } #[test] @@ -176,20 +174,18 @@ pub fn basic_test_2() { ); let batch = reader.collect::, _>>().unwrap(); - let expected = r#"+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+ -| int_short_repeated | int_neg_short_repeated | int_delta | int_neg_delta | int_direct | int_neg_direct | bigint_direct | bigint_neg_direct | bigint_other | utf8_increase | utf8_decrease | -+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+ -| 5 | -5 | 1 | 5 | 1 | -1 | 1 | -1 | 5 | a | eeeee | -| 5 | -5 | 2 | 4 | 6 | -6 | 6 | -6 | -5 | bb | dddd | -| | | | | | | | | 1 | ccc | ccc | -| 5 | -5 | 4 | 2 | 3 | -3 | 3 | -3 | 5 | dddd | bb | -| 5 | -5 | 5 | 1 | 2 | -2 | 2 | -2 | 5 | eeeee | a | -+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+"#; - - assert_eq!( - expected, - pretty::pretty_format_batches(&batch).unwrap().to_string() - ) + let expected = [ + "+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+", + "| int_short_repeated | int_neg_short_repeated | int_delta | int_neg_delta | int_direct | int_neg_direct | bigint_direct | bigint_neg_direct | bigint_other | utf8_increase | utf8_decrease |", + "+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+", + "| 5 | -5 | 1 | 5 | 1 | -1 | 1 | -1 | 5 | a | eeeee |", + "| 5 | -5 | 2 | 4 | 6 | -6 | 6 | -6 | -5 | bb | dddd |", + "| | | | | | | | | 1 | ccc | ccc |", + "| 5 | -5 | 4 | 2 | 3 | -3 | 3 | -3 | 5 | dddd | bb |", + "| 5 | -5 | 5 | 1 | 2 | -2 | 2 | -2 | 5 | eeeee | a |", + "+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+", + ]; + assert_batches_eq(&batch, &expected); } #[test] @@ -198,19 +194,18 @@ pub fn basic_test_3() { let reader = new_arrow_reader(&path, &["timestamp_simple", "date_simple"]); let batch = reader.collect::, _>>().unwrap(); - let expected = r#"+----------------------------+-------------+ -| timestamp_simple | date_simple | -+----------------------------+-------------+ -| 2023-04-01T20:15:30.002 | 2023-04-01 | -| 2021-08-22T07:26:44.525777 | 2023-03-01 | -| 2023-01-01T00:00:00 | 2023-01-01 | -| 2023-02-01T00:00:00 | 2023-02-01 | -| 2023-03-01T00:00:00 | 2023-03-01 | -+----------------------------+-------------+"#; - assert_eq!( - expected, - pretty::pretty_format_batches(&batch).unwrap().to_string() - ) + let expected = [ + "+----------------------------+-------------+", + "| timestamp_simple | date_simple |", + "+----------------------------+-------------+", + "| 2023-04-01T20:15:30.002 | 2023-04-01 |", + "| 2021-08-22T07:26:44.525777 | 2023-03-01 |", + "| 2023-01-01T00:00:00 | 2023-01-01 |", + "| 2023-02-01T00:00:00 | 2023-02-01 |", + "| 2023-03-01T00:00:00 | 2023-03-01 |", + "+----------------------------+-------------+", + ]; + assert_batches_eq(&batch, &expected); } #[test] @@ -218,19 +213,18 @@ pub fn basic_test_nested_struct() { let path = basic_path("nested_struct.orc"); let reader = new_arrow_reader_root(&path); let batch = reader.collect::, _>>().unwrap(); - let expected = r#"+-------------------+ -| nest | -+-------------------+ -| {a: 1.0, b: true} | -| {a: 3.0, b: } | -| {a: , b: } | -| | -| {a: -3.0, b: } | -+-------------------+"#; - assert_eq!( - expected, - pretty::pretty_format_batches(&batch).unwrap().to_string() - ) + let expected = [ + "+-------------------+", + "| nest |", + "+-------------------+", + "| {a: 1.0, b: true} |", + "| {a: 3.0, b: } |", + "| {a: , b: } |", + "| |", + "| {a: -3.0, b: } |", + "+-------------------+", + ]; + assert_batches_eq(&batch, &expected); } #[test] @@ -239,19 +233,18 @@ pub fn basic_test_nested_array() { let reader = new_arrow_reader_root(&path); let batch = reader.collect::, _>>().unwrap(); - let expected = r#"+--------------------+ -| value | -+--------------------+ -| [1, , 3, 43, 5] | -| [5, , 32, 4, 15] | -| [16, , 3, 4, 5, 6] | -| | -| [3, ] | -+--------------------+"#; - assert_eq!( - expected, - pretty::pretty_format_batches(&batch).unwrap().to_string() - ) + let expected = [ + "+--------------------+", + "| value |", + "+--------------------+", + "| [1, , 3, 43, 5] |", + "| [5, , 32, 4, 15] |", + "| [16, , 3, 4, 5, 6] |", + "| |", + "| [3, ] |", + "+--------------------+", + ]; + assert_batches_eq(&batch, &expected); } #[test] @@ -260,18 +253,17 @@ pub fn basic_test_nested_map() { let reader = new_arrow_reader_root(&path); let batch = reader.collect::, _>>().unwrap(); - let expected = r#"+--------------------------+ -| map | -+--------------------------+ -| {zero: 0, one: 1} | -| | -| {two: 2, tree: 3} | -| {one: 1, two: 2, nill: } | -+--------------------------+"#; - assert_eq!( - expected, - pretty::pretty_format_batches(&batch).unwrap().to_string() - ) + let expected = [ + "+--------------------------+", + "| map |", + "+--------------------------+", + "| {zero: 0, one: 1} |", + "| |", + "| {two: 2, tree: 3} |", + "| {one: 1, two: 2, nill: } |", + "+--------------------------+", + ]; + assert_batches_eq(&batch, &expected); } #[test] @@ -280,19 +272,18 @@ pub fn basic_test_0() { let reader = new_arrow_reader_root(&path); let batch = reader.collect::, _>>().unwrap(); - let expected = r#"+-----+-------+------------+-----+-----+-------+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+----------------------------+-------------+----------------+ -| a | b | str_direct | d | e | f | int_short_repeated | int_neg_short_repeated | int_delta | int_neg_delta | int_direct | int_neg_direct | bigint_direct | bigint_neg_direct | bigint_other | utf8_increase | utf8_decrease | timestamp_simple | date_simple | tinyint_simple | -+-----+-------+------------+-----+-----+-------+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+----------------------------+-------------+----------------+ -| 1.0 | true | a | a | ddd | aaaaa | 5 | -5 | 1 | 5 | 1 | -1 | 1 | -1 | 5 | a | eeeee | 2023-04-01T20:15:30.002 | 2023-04-01 | -1 | -| 2.0 | false | cccccc | bb | cc | bbbbb | 5 | -5 | 2 | 4 | 6 | -6 | 6 | -6 | -5 | bb | dddd | 2021-08-22T07:26:44.525777 | 2023-03-01 | | -| | | | | | | | | | | | | | | 1 | ccc | ccc | 2023-01-01T00:00:00 | 2023-01-01 | 1 | -| 4.0 | true | ddd | ccc | bb | ccccc | 5 | -5 | 4 | 2 | 3 | -3 | 3 | -3 | 5 | dddd | bb | 2023-02-01T00:00:00 | 2023-02-01 | 127 | -| 5.0 | false | ee | ddd | a | ddddd | 5 | -5 | 5 | 1 | 2 | -2 | 2 | -2 | 5 | eeeee | a | 2023-03-01T00:00:00 | 2023-03-01 | -127 | -+-----+-------+------------+-----+-----+-------+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+----------------------------+-------------+----------------+"#; - assert_eq!( - expected, - pretty::pretty_format_batches(&batch).unwrap().to_string() - ) + let expected = [ + "+-----+-------+------------+-----+-----+-------+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+----------------------------+-------------+----------------+", + "| a | b | str_direct | d | e | f | int_short_repeated | int_neg_short_repeated | int_delta | int_neg_delta | int_direct | int_neg_direct | bigint_direct | bigint_neg_direct | bigint_other | utf8_increase | utf8_decrease | timestamp_simple | date_simple | tinyint_simple |", + "+-----+-------+------------+-----+-----+-------+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+----------------------------+-------------+----------------+", + "| 1.0 | true | a | a | ddd | aaaaa | 5 | -5 | 1 | 5 | 1 | -1 | 1 | -1 | 5 | a | eeeee | 2023-04-01T20:15:30.002 | 2023-04-01 | -1 |", + "| 2.0 | false | cccccc | bb | cc | bbbbb | 5 | -5 | 2 | 4 | 6 | -6 | 6 | -6 | -5 | bb | dddd | 2021-08-22T07:26:44.525777 | 2023-03-01 | |", + "| | | | | | | | | | | | | | | 1 | ccc | ccc | 2023-01-01T00:00:00 | 2023-01-01 | 1 |", + "| 4.0 | true | ddd | ccc | bb | ccccc | 5 | -5 | 4 | 2 | 3 | -3 | 3 | -3 | 5 | dddd | bb | 2023-02-01T00:00:00 | 2023-02-01 | 127 |", + "| 5.0 | false | ee | ddd | a | ddddd | 5 | -5 | 5 | 1 | 2 | -2 | 2 | -2 | 5 | eeeee | a | 2023-03-01T00:00:00 | 2023-03-01 | -127 |", + "+-----+-------+------------+-----+-----+-------+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+----------------------------+-------------+----------------+", + ]; + assert_batches_eq(&batch, &expected); } #[tokio::test] @@ -301,19 +292,18 @@ pub async fn async_basic_test_0() { let reader = new_arrow_stream_reader_root(&path).await; let batch = reader.try_collect::>().await.unwrap(); - let expected = r#"+-----+-------+------------+-----+-----+-------+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+----------------------------+-------------+----------------+ -| a | b | str_direct | d | e | f | int_short_repeated | int_neg_short_repeated | int_delta | int_neg_delta | int_direct | int_neg_direct | bigint_direct | bigint_neg_direct | bigint_other | utf8_increase | utf8_decrease | timestamp_simple | date_simple | tinyint_simple | -+-----+-------+------------+-----+-----+-------+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+----------------------------+-------------+----------------+ -| 1.0 | true | a | a | ddd | aaaaa | 5 | -5 | 1 | 5 | 1 | -1 | 1 | -1 | 5 | a | eeeee | 2023-04-01T20:15:30.002 | 2023-04-01 | -1 | -| 2.0 | false | cccccc | bb | cc | bbbbb | 5 | -5 | 2 | 4 | 6 | -6 | 6 | -6 | -5 | bb | dddd | 2021-08-22T07:26:44.525777 | 2023-03-01 | | -| | | | | | | | | | | | | | | 1 | ccc | ccc | 2023-01-01T00:00:00 | 2023-01-01 | 1 | -| 4.0 | true | ddd | ccc | bb | ccccc | 5 | -5 | 4 | 2 | 3 | -3 | 3 | -3 | 5 | dddd | bb | 2023-02-01T00:00:00 | 2023-02-01 | 127 | -| 5.0 | false | ee | ddd | a | ddddd | 5 | -5 | 5 | 1 | 2 | -2 | 2 | -2 | 5 | eeeee | a | 2023-03-01T00:00:00 | 2023-03-01 | -127 | -+-----+-------+------------+-----+-----+-------+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+----------------------------+-------------+----------------+"#; - assert_eq!( - expected, - pretty::pretty_format_batches(&batch).unwrap().to_string() - ) + let expected = [ + "+-----+-------+------------+-----+-----+-------+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+----------------------------+-------------+----------------+", + "| a | b | str_direct | d | e | f | int_short_repeated | int_neg_short_repeated | int_delta | int_neg_delta | int_direct | int_neg_direct | bigint_direct | bigint_neg_direct | bigint_other | utf8_increase | utf8_decrease | timestamp_simple | date_simple | tinyint_simple |", + "+-----+-------+------------+-----+-----+-------+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+----------------------------+-------------+----------------+", + "| 1.0 | true | a | a | ddd | aaaaa | 5 | -5 | 1 | 5 | 1 | -1 | 1 | -1 | 5 | a | eeeee | 2023-04-01T20:15:30.002 | 2023-04-01 | -1 |", + "| 2.0 | false | cccccc | bb | cc | bbbbb | 5 | -5 | 2 | 4 | 6 | -6 | 6 | -6 | -5 | bb | dddd | 2021-08-22T07:26:44.525777 | 2023-03-01 | |", + "| | | | | | | | | | | | | | | 1 | ccc | ccc | 2023-01-01T00:00:00 | 2023-01-01 | 1 |", + "| 4.0 | true | ddd | ccc | bb | ccccc | 5 | -5 | 4 | 2 | 3 | -3 | 3 | -3 | 5 | dddd | bb | 2023-02-01T00:00:00 | 2023-02-01 | 127 |", + "| 5.0 | false | ee | ddd | a | ddddd | 5 | -5 | 5 | 1 | 2 | -2 | 2 | -2 | 5 | eeeee | a | 2023-03-01T00:00:00 | 2023-03-01 | -127 |", + "+-----+-------+------------+-----+-----+-------+--------------------+------------------------+-----------+---------------+------------+----------------+---------------+-------------------+--------------+---------------+---------------+----------------------------+-------------+----------------+", + ]; + assert_batches_eq(&batch, &expected); } #[test]