From 0290cdb5d659f22af0a4481cf185549ff9563a1e Mon Sep 17 00:00:00 2001
From: evenyag <realevenyag@gmail.com>
Date: Fri, 9 Sep 2022 16:10:57 +0800
Subject: [PATCH] test: Fix merge tests (#243)

* test: Fix merge tests

The merge tests still use Int64Vector for timestamp, which should
use TimestampVector instead.

* test: Test Debug format for Source::Reader

Mainly for improve code coverage
---
 src/storage/src/read/merge.rs          |  8 +++++---
 src/storage/src/test_util/read_util.rs | 12 ++++++------
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/src/storage/src/read/merge.rs b/src/storage/src/read/merge.rs
index 80ce3f74fa..a1f69b2e34 100644
--- a/src/storage/src/read/merge.rs
+++ b/src/storage/src/read/merge.rs
@@ -591,7 +591,7 @@ impl MergeReader {
 #[cfg(test)]
 mod tests {
     use datatypes::prelude::ScalarVector;
-    use datatypes::vectors::Int64Vector;
+    use datatypes::vectors::{Int64Vector, TimestampVector};
 
     use super::*;
     use crate::test_util::read_util;
@@ -633,6 +633,8 @@ mod tests {
         let output = format!("{:?}", left);
         assert!(output.contains("cursor"));
         assert!(output.contains("pos: 1"));
+        let output = format!("{:?}", right);
+        assert!(output.contains("cursor"));
         let output = format!("{:?}", left.first_row());
         assert!(output.contains("pos: 1"));
     }
@@ -676,7 +678,7 @@ mod tests {
             let key = batch
                 .column(0)
                 .as_any()
-                .downcast_ref::<Int64Vector>()
+                .downcast_ref::<TimestampVector>()
                 .unwrap();
             let value = batch
                 .column(1)
@@ -687,7 +689,7 @@ mod tests {
             let batch: Vec<_> = key
                 .iter_data()
                 .zip(value.iter_data())
-                .map(|(k, v)| (k.unwrap(), v))
+                .map(|(k, v)| (k.unwrap().value(), v))
                 .collect();
             result.push(batch);
         }
diff --git a/src/storage/src/test_util/read_util.rs b/src/storage/src/test_util/read_util.rs
index 4570a941c8..d18868682b 100644
--- a/src/storage/src/test_util/read_util.rs
+++ b/src/storage/src/test_util/read_util.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 use async_trait::async_trait;
 use datatypes::prelude::ScalarVector;
 use datatypes::type_id::LogicalTypeId;
-use datatypes::vectors::{Int64Vector, UInt64Vector, UInt8Vector};
+use datatypes::vectors::{Int64Vector, TimestampVector, UInt64Vector, UInt8Vector};
 
 use crate::error::Result;
 use crate::memtable::{BatchIterator, BoxedBatchIterator, RowOrdering};
@@ -30,7 +30,7 @@ pub fn new_projected_schema() -> ProjectedSchemaRef {
 
 /// Build a new batch, with 0 sequence and op_type.
 fn new_kv_batch(key_values: &[(i64, Option<i64>)]) -> Batch {
-    let key = Arc::new(Int64Vector::from_values(key_values.iter().map(|v| v.0)));
+    let key = Arc::new(TimestampVector::from_values(key_values.iter().map(|v| v.0)));
     let value = Arc::new(Int64Vector::from_iter(key_values.iter().map(|v| v.1)));
     let sequences = Arc::new(UInt64Vector::from_vec(vec![0; key_values.len()]));
     let op_types = Arc::new(UInt8Vector::from_vec(vec![0; key_values.len()]));
@@ -43,7 +43,7 @@ fn check_kv_batch(batches: &[Batch], expect: &[&[(i64, Option<i64>)]]) {
         let key = batch
             .column(0)
             .as_any()
-            .downcast_ref::<Int64Vector>()
+            .downcast_ref::<TimestampVector>()
             .unwrap();
         let value = batch
             .column(1)
@@ -52,7 +52,7 @@ fn check_kv_batch(batches: &[Batch], expect: &[&[(i64, Option<i64>)]]) {
             .unwrap();
 
         for (i, (k, v)) in key_values.iter().enumerate() {
-            assert_eq!(key.get_data(i).unwrap(), *k,);
+            assert_eq!(key.get_data(i).unwrap().value(), *k);
             assert_eq!(value.get_data(i), *v,);
         }
     }
@@ -65,7 +65,7 @@ pub async fn collect_kv_batch(reader: &mut dyn BatchReader) -> Vec<(i64, Option<
         let key = batch
             .column(0)
             .as_any()
-            .downcast_ref::<Int64Vector>()
+            .downcast_ref::<TimestampVector>()
             .unwrap();
         let value = batch
             .column(1)
@@ -74,7 +74,7 @@ pub async fn collect_kv_batch(reader: &mut dyn BatchReader) -> Vec<(i64, Option<
             .unwrap();
 
         for (k, v) in key.iter_data().zip(value.iter_data()) {
-            result.push((k.unwrap(), v));
+            result.push((k.unwrap().value(), v));
         }
     }