From 32b781a521ca99dea63d521176bcfbc2451ce48d Mon Sep 17 00:00:00 2001
From: "Jorge C. Leitao" <jorgecarleitao@gmail.com>
Date: Wed, 22 Jun 2022 05:18:15 +0000
Subject: [PATCH] Added more tests

---
 .github/workflows/test.yml |  2 +-
 src/io/flight/mod.rs       |  6 ++++--
 tests/it/io/flight/mod.rs  | 32 ++++++++++++++++++++++++++++++++
 tests/it/io/mod.rs         |  3 +++
 4 files changed, 40 insertions(+), 3 deletions(-)
 create mode 100644 tests/it/io/flight/mod.rs

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 58abc32ef93..40015100d00 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -88,7 +88,7 @@ jobs:
 
       - name: Run
         # --skip io: miri can't handle opening of files, so we skip those
-        run: cargo miri test --features full -- --skip io::parquet --skip io::ipc
+        run: cargo miri test --features full -- --skip io::parquet --skip io::ipc --skip io::flight
 
   miri-checks-io:
     name: MIRI on IO IPC
diff --git a/src/io/flight/mod.rs b/src/io/flight/mod.rs
index b439ced2d9b..8aa4a1de0e1 100644
--- a/src/io/flight/mod.rs
+++ b/src/io/flight/mod.rs
@@ -11,12 +11,14 @@ use crate::{
     error::{Error, Result},
     io::ipc::read,
     io::ipc::write,
-    io::ipc::write::common::{encode_chunk, DictionaryTracker, EncodedData, WriteOptions},
+    io::ipc::write::common::{encode_chunk, DictionaryTracker, EncodedData},
 };
 
-pub use super::ipc::write::default_ipc_fields;
 use super::ipc::{IpcField, IpcSchema};
 
+pub use super::ipc::write::default_ipc_fields;
+pub use crate::io::ipc::write::common::WriteOptions;
+
 /// Serializes [`Chunk`] to a vector of [`FlightData`] representing the serialized dictionaries
 /// and a [`FlightData`] representing the batch.
 /// # Errors
diff --git a/tests/it/io/flight/mod.rs b/tests/it/io/flight/mod.rs
new file mode 100644
index 00000000000..9718b749027
--- /dev/null
+++ b/tests/it/io/flight/mod.rs
@@ -0,0 +1,32 @@
+use arrow2::array::Array;
+use arrow2::chunk::Chunk;
+use arrow2::datatypes::Schema;
+use arrow2::error::Error;
+
+use arrow2::io::flight::*;
+use arrow2::io::ipc::write::{default_ipc_fields, WriteOptions};
+
+use super::ipc::read_gzip_json;
+
+fn round_trip(schema: Schema, chunk: Chunk<Box<dyn Array>>) -> Result<(), Error> {
+    let fields = default_ipc_fields(&schema.fields);
+    let serialized = serialize_schema(&schema, Some(&fields));
+    let (result, ipc_schema) = deserialize_schemas(&serialized.data_header)?;
+    assert_eq!(schema, result);
+
+    let (_, batch) = serialize_batch(&chunk, &fields, &WriteOptions { compression: None })?;
+
+    let result = deserialize_batch(&batch, &result.fields, &ipc_schema, &Default::default())?;
+    assert_eq!(result, chunk);
+    Ok(())
+}
+
+#[test]
+fn generated_nested_dictionary() -> Result<(), Error> {
+    let (schema, _, mut batches) =
+        read_gzip_json("1.0.0-littleendian", "generated_nested").unwrap();
+
+    round_trip(schema, batches.pop().unwrap())?;
+
+    Ok(())
+}
diff --git a/tests/it/io/mod.rs b/tests/it/io/mod.rs
index d192b806ee4..3fc0c2123e3 100644
--- a/tests/it/io/mod.rs
+++ b/tests/it/io/mod.rs
@@ -22,3 +22,6 @@ mod avro;
     feature = "io_csv_read_async"
 ))]
 mod csv;
+
+#[cfg(feature = "io_flight")]
+mod flight;