add additional tests to cover with_state cases

Author: milenkovicm

ballista/client/tests/common/mod.rs

@@ -25,7 +25,7 @@ use ballista_core::serde::{
 };
 use ballista_core::{ConfigProducer, RuntimeProducer};
 use ballista_scheduler::SessionBuilder;
-use datafusion::execution::SessionState;
+use datafusion::execution::{SessionState, SessionStateBuilder};
 use datafusion::prelude::{SessionConfig, SessionContext};
 
 /// Returns the parquet test data directory, which is by default
@@ -234,6 +234,21 @@ pub async fn remote_context() -> SessionContext {
         .unwrap()
 }
 
+#[allow(dead_code)]
+pub async fn standalone_context_with_state() -> SessionContext {
+    let state = SessionStateBuilder::new().with_default_features().build();
+    SessionContext::standalone_with_state(state).await.unwrap()
+}
+
+#[allow(dead_code)]
+pub async fn remote_context_with_state() -> SessionContext {
+    let state = SessionStateBuilder::new().with_default_features().build();
+    let (host, port) = setup_test_cluster_with_state(state.clone()).await;
+    SessionContext::remote_with_state(&format!("df://{host}:{port}"), state)
+        .await
+        .unwrap()
+}
+
 #[ctor::ctor]
 fn init() {
     // Enable RUST_LOG logging configuration for test

ballista/client/tests/context_checks.rs

@@ -19,7 +19,10 @@ mod common;
 #[cfg(test)]
 mod supported {
 
-    use crate::common::{remote_context, standalone_context};
+    use crate::common::{
+        remote_context, remote_context_with_state, standalone_context,
+        standalone_context_with_state,
+    };
     use ballista_core::config::BallistaConfig;
     use datafusion::prelude::*;
     use datafusion::{assert_batches_eq, prelude::SessionContext};
@@ -499,6 +502,8 @@ mod supported {
     #[rstest]
     #[case::standalone(standalone_context())]
     #[case::remote(remote_context())]
+    #[case::standalone_state(standalone_context_with_state())]
+    #[case::remote_state(remote_context_with_state())]
     #[tokio::test]
     #[cfg(not(windows))] // test is failing at windows, can't debug it
     async fn should_execute_sql_write_read_roundtrip(
@@ -587,6 +592,8 @@ mod supported {
     #[rstest]
     #[case::standalone(standalone_context())]
     #[case::remote(remote_context())]
+    #[case::standalone_state(standalone_context_with_state())]
+    #[case::remote_state(remote_context_with_state())]
     #[tokio::test]
     async fn should_execute_sql_show_multiple_times(
         #[future(awt)]