refactor tests with insta

Author: dmitriibugakov

.gitignore

@@ -1,69 +1,59 @@
+# License
 # Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements.  See the NOTICE file
+# or more contributor license agreements. See the NOTICE file
 # distributed with this work for additional information
-# regarding copyright ownership.  The ASF licenses this file
+# regarding copyright ownership. The ASF licenses this file
 # to you under the Apache License, Version 2.0 (the
 # "License"); you may not use this file except in compliance
-# with the License.  You may obtain a copy of the License at
+# with the License. You may obtain a copy of the License at
 #
-#   http://www.apache.org/licenses/LICENSE-2.0
+# http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied.  See the License for the
+# KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations
 # under the License.
 
-apache-rat-*.jar
-
 # Linux perf sample data
 perf.data
 perf.data.old
 
-
+# IDE and editor directories
 .vscode
 .idea/
 .pytest_cache/
-pkgs
-docker_cache
-.gdb_history
-*.orig
-.*.swp
-.*.swo
 
+# Python virtual environments
 venv/*
 
 # macOS
 .DS_Store
 
-# docker volumes used for caching
+# Docker volumes used for caching
 .docker
 
-# Rust
+# Rust build output
 target
 Cargo.lock
-!datafusion-cli/Cargo.lock
 
-rusty-tags.vi
+# Backup files
+*.orig
+.gdb_history
 .history
-.flatbuffers/
-
-# apache release artifacts
-dev/dist
-
-# CI
-arrow-ballista
 
-datafusion/CHANGELOG.md.bak
-.githubchangeloggenerator.cache*
+# Swap files
+.*.swp
+.*.swo
 
-# Generated tpch data
-datafusion/sqllogictests/test_files/tpch/data/*
+# Rust-specific tools
+rusty-tags.vi
 
-# Scratch temp dir for sqllogictests
-datafusion/sqllogictest/test_files/scratch*
+# Apache release artifacts
+dev/dist
 
-# rat
+# Apache RAT report files
 filtered_rat.txt
 rat.txt
+.githubchangeloggenerator.cache*

Cargo.toml

@@ -42,6 +42,7 @@ paste = "1"
 [dev-dependencies]
 arrow = { version = "53.0.0", features = ["test_utils"] }
 criterion = { version = "0.5", features = ["async_tokio"] }
+insta = { version = "1.40.0", features = ["yaml"] }
 tokio = { version = "1.36", features = ["full"] }
 
 [lints.clippy]

README.md

@@ -1,7 +1,7 @@
 # datafusion-functions-extra
 
-<!-- [![CI](https://github.com/datafusion-contrib/datafusion-functions-json/actions/workflows/ci.yml/badge.svg?event=push)](https://github.com/datafusion-contrib/datafusion-functions-json/actions/workflows/ci.yml?query=branch%3Amain) -->
-<!-- [![Crates.io](https://img.shields.io/crates/v/datafusion-functions-json?color=green)](https://crates.io/crates/datafusion-functions-json) -->
+[![CI](https://github.com/datafusion-contrib/datafusion-functions-extra/actions/workflows/ci.yml/badge.svg?event=push)](https://github.com/datafusion-contrib/datafusion-functions-extra/actions/workflows/ci.yml?query=branch%3Amain)
+<!-- [![Crates.io](https://img.shields.io/crates/v/datafusion-functions-extra?color=green)](https://crates.io/crates/datafusion-functions-extra) -->
 
 **Note:** This is not an official Apache Software Foundation release.
 

tests/main.rs

@@ -15,48 +15,103 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::datatypes::{DataType, TimeUnit};
-use utils::{display_val, run_query};
+use crate::utils::TestExecution;
 
 mod utils;
 
+static TEST_TABLE: &str = r#"
+CREATE TABLE test_table (
+    utf8_col VARCHAR,
+    int64_col BIGINT,
+    float64_col DOUBLE,
+    date64_col DATE,
+    time64_col TIME
+) AS VALUES
+    ('apple', 1, 1.0, DATE '2021-01-01', TIME '01:00:00'),
+    ('banana', 2, 2.0, DATE '2021-01-02', TIME '02:00:00'),
+    ('apple', 2, 2.0, DATE '2021-01-02', TIME '02:00:00'),
+    ('orange', 3, 3.0, DATE '2021-01-03', TIME '03:00:00'),
+    ('banana', 3, 3.0, DATE '2021-01-03', TIME '03:00:00'),
+    ('apple', 3, 3.0, DATE '2021-01-03', TIME '03:00:00'),
+    (NULL, NULL, NULL, NULL, NULL);
+"#;
+
 #[tokio::test]
 async fn test_mode_utf8() {
-    let sql = "SELECT MODE(utf8_col) FROM test_table";
-    let batches = run_query(sql).await.unwrap();
-    assert_eq!(display_val(batches).await, (DataType::Utf8, "apple".to_string()));
-}
+    let mut execution = TestExecution::new().await.unwrap().with_setup(TEST_TABLE).await;
+
+    let actual = execution.run_and_format("SELECT MODE(utf8_col) FROM test_table").await;
 
+    insta::assert_yaml_snapshot!(actual, @r###"
+          - +---------------------------+
+          - "| mode(test_table.utf8_col) |"
+          - +---------------------------+
+          - "| apple                     |"
+          - +---------------------------+
+    "###);
+}
 #[tokio::test]
 async fn test_mode_int64() {
-    let sql = "SELECT MODE(int64_col) FROM test_table";
-    let batches = run_query(sql).await.unwrap();
-    assert_eq!(display_val(batches).await, (DataType::Int64, "3".to_string()));
+    let mut execution = TestExecution::new().await.unwrap().with_setup(TEST_TABLE).await;
+
+    let actual = execution.run_and_format("SELECT MODE(int64_col) FROM test_table").await;
+
+    insta::assert_yaml_snapshot!(actual, @r###"
+          - +----------------------------+
+          - "| mode(test_table.int64_col) |"
+          - +----------------------------+
+          - "| 3                          |"
+          - +----------------------------+
+    "###);
 }
 
 #[tokio::test]
 async fn test_mode_float64() {
-    let sql = "SELECT MODE(float64_col) FROM test_table";
-    let batches = run_query(sql).await.unwrap();
-    assert_eq!(display_val(batches).await, (DataType::Float64, "3.0".to_string()));
+    let mut execution = TestExecution::new().await.unwrap().with_setup(TEST_TABLE).await;
+
+    let actual = execution
+        .run_and_format("SELECT MODE(float64_col) FROM test_table")
+        .await;
+
+    insta::assert_yaml_snapshot!(actual, @r###"
+          - +------------------------------+
+          - "| mode(test_table.float64_col) |"
+          - +------------------------------+
+          - "| 3.0                          |"
+          - +------------------------------+
+    "###);
 }
 
 #[tokio::test]
 async fn test_mode_date64() {
-    let sql = "SELECT MODE(date64_col) FROM test_table";
-    let batches = run_query(sql).await.unwrap();
-    assert_eq!(
-        display_val(batches).await,
-        (DataType::Date64, "2021-01-03T00:00:00".to_string())
-    );
+    let mut execution = TestExecution::new().await.unwrap().with_setup(TEST_TABLE).await;
+
+    let actual = execution
+        .run_and_format("SELECT MODE(date64_col) FROM test_table")
+        .await;
+
+    insta::assert_yaml_snapshot!(actual, @r###"
+          - +-----------------------------+
+          - "| mode(test_table.date64_col) |"
+          - +-----------------------------+
+          - "| 2021-01-03                  |"
+          - +-----------------------------+
+    "###);
 }
 
 #[tokio::test]
 async fn test_mode_time64() {
-    let sql = "SELECT MODE(time64_col) FROM test_table";
-    let batches = run_query(sql).await.unwrap();
-    assert_eq!(
-        display_val(batches).await,
-        (DataType::Time64(TimeUnit::Microsecond), "03:00:00".to_string())
-    );
+    let mut execution = TestExecution::new().await.unwrap().with_setup(TEST_TABLE).await;
+
+    let actual = execution
+        .run_and_format("SELECT MODE(time64_col) FROM test_table")
+        .await;
+
+    insta::assert_yaml_snapshot!(actual, @r###"
+          - +-----------------------------+
+          - "| mode(test_table.time64_col) |"
+          - +-----------------------------+
+          - "| 03:00:00                    |"
+          - +-----------------------------+
+    "###);
 }

tests/utils/mod.rs

@@ -15,113 +15,57 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::sync::Arc;
-
-use arrow::array::{Date64Array, Float64Array, Time64MicrosecondArray};
-use arrow::datatypes::TimeUnit;
-use datafusion::arrow::array::{ArrayRef, Int64Array};
-use datafusion::arrow::datatypes::{DataType, Field, Schema};
-use datafusion::arrow::util::display::{ArrayFormatter, FormatOptions};
-use datafusion::arrow::{array::StringArray, record_batch::RecordBatch};
+use arrow::record_batch::RecordBatch;
+use arrow::util::pretty::pretty_format_batches;
 use datafusion::error::Result;
 use datafusion::execution::context::SessionContext;
 use datafusion::prelude::SessionConfig;
+use datafusion::sql::parser::DFParser;
 use datafusion_functions_extra::register_all_extra_functions;
+use log::debug;
 
-// TODO: It would be great to release `datafusion-sqllogictest` as a crate.
-// This would allow easy integration and testing of SQL queries across projects.
-pub async fn create_context() -> Result<SessionContext> {
-    let config = SessionConfig::new();
-    let mut ctx = SessionContext::new_with_config(config);
-    register_all_extra_functions(&mut ctx)?;
-    Ok(ctx)
+pub struct TestExecution {
+    ctx: SessionContext,
 }
 
-async fn create_test_table() -> Result<SessionContext> {
-    let ctx = create_context().await?;
-
-    let schema = Arc::new(Schema::new(vec![
-        Field::new("utf8_col", DataType::Utf8, true),
-        Field::new("int64_col", DataType::Int64, true),
-        Field::new("float64_col", DataType::Float64, true),
-        Field::new("date64_col", DataType::Date64, true),
-        Field::new("time64_col", DataType::Time64(TimeUnit::Microsecond), true),
-    ]));
-
-    let utf8_values: ArrayRef = Arc::new(StringArray::from(vec![
-        Some("apple"),
-        Some("banana"),
-        Some("apple"),
-        Some("orange"),
-        Some("banana"),
-        Some("apple"),
-        None,
-    ]));
-
-    let int64_values: ArrayRef = Arc::new(Int64Array::from(vec![
-        Some(1),
-        Some(2),
-        Some(2),
-        Some(3),
-        Some(3),
-        Some(3),
-        None,
-    ]));
+impl TestExecution {
+    pub async fn new() -> Result<Self> {
+        let config = SessionConfig::new();
+        let mut ctx = SessionContext::new_with_config(config);
+        register_all_extra_functions(&mut ctx)?;
+        Ok(Self { ctx })
+    }
 
-    let float64_values: ArrayRef = Arc::new(Float64Array::from(vec![
-        Some(1.0),
-        Some(2.0),
-        Some(2.0),
-        Some(3.0),
-        Some(3.0),
-        Some(3.0),
-        None,
-    ]));
+    pub async fn with_setup(self, sql: &str) -> Self {
+        debug!("Running setup query: {sql}");
+        let statements = DFParser::parse_sql(sql).expect("Error parsing setup query");
+        for statement in statements {
+            debug!("Running setup statement: {statement}");
+            let statement_sql = statement.to_string();
+            self.ctx
+                .sql(&statement_sql)
+                .await
+                .expect("Error planning setup failed")
+                .collect()
+                .await
+                .expect("Error executing setup query");
+        }
+        self
+    }
 
-    let date64_values: ArrayRef = Arc::new(Date64Array::from(vec![
-        Some(1609459200000),
-        Some(1609545600000),
-        Some(1609545600000),
-        Some(1609632000000),
-        Some(1609632000000),
-        Some(1609632000000),
-        None,
-    ]));
+    pub async fn run(&mut self, sql: &str) -> Result<Vec<RecordBatch>> {
+        debug!("Running query: {sql}");
+        self.ctx.sql(sql).await?.collect().await
+    }
 
-    let time64_values: ArrayRef = Arc::new(Time64MicrosecondArray::from(vec![
-        Some(3600000000),
-        Some(7200000000),
-        Some(7200000000),
-        Some(10800000000),
-        Some(10800000000),
-        Some(10800000000),
-        None,
-    ]));
-
-    let batch = RecordBatch::try_new(
-        schema.clone(),
-        vec![utf8_values, int64_values, float64_values, date64_values, time64_values],
-    )?;
-
-    ctx.register_batch("test_table", batch)?;
-
-    Ok(ctx)
+    pub async fn run_and_format(&mut self, sql: &str) -> Vec<String> {
+        let results = self.run(sql).await.expect("Error running query");
+        format_results(&results)
+    }
 }
 
-pub async fn run_query(sql: &str) -> Result<Vec<RecordBatch>> {
-    let ctx = create_test_table().await?;
-    ctx.sql(sql).await?.collect().await
-}
+fn format_results(results: &[RecordBatch]) -> Vec<String> {
+    let formatted = pretty_format_batches(results).unwrap().to_string();
 
-pub async fn display_val(batch: Vec<RecordBatch>) -> (DataType, String) {
-    assert_eq!(batch.len(), 1);
-    let batch = batch.first().unwrap();
-    assert_eq!(batch.num_rows(), 1);
-    let schema = batch.schema();
-    let schema_col = schema.field(0);
-    let c = batch.column(0);
-    let options = FormatOptions::default().with_display_error(true);
-    let f = ArrayFormatter::try_new(c.as_ref(), &options).unwrap();
-    let repr = f.value(0).try_to_string().unwrap();
-    (schema_col.data_type().clone(), repr)
+    formatted.lines().map(|s| s.to_string()).collect()
 }