Adjust Python operator API tests

dora-rs · Oct 19, 2023 · df65eab · df65eab
1 parent 6140681
commit df65eab
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/apis/python/operator/src/lib.rs b/apis/python/operator/src/lib.rs
@@ -160,18 +160,19 @@ mod tests {
     };
 
     use arrow_schema::{DataType, Field};
-    use dora_node_api::Data;
+    use dora_node_api::{
+        arrow_utils::{copy_array_into_sample, required_data_size},
+        RawData,
+    };
     use eyre::{Context, Result};
 
-    use crate::{copy_array_into_sample, required_data_size};
-
     fn assert_roundtrip(arrow_array: &ArrayData) -> Result<()> {
         let size = required_data_size(&arrow_array);
         let mut sample: Vec<u8> = vec![0; size];
 
         let info = copy_array_into_sample(&mut sample, &arrow_array)?;
 
-        let serialized_deserialized_arrow_array = Arc::new(Data::Vec(sample))
+        let serialized_deserialized_arrow_array = Arc::new(RawData::Vec(sample))
             .into_arrow_array(&info)
             .context("Could not create arrow array")?;
         assert_eq!(arrow_array, &serialized_deserialized_arrow_array);