add config lookup and force connect logic ~cut-release

chris-twiner · chris-twiner · commit d4fdd09866e6 · 2025-12-10T11:06:48.000+01:00
diff --git a/runtime/src/main/scala/com/sparkutils/testing/TestUtilsEnvironment.scala b/runtime/src/main/scala/com/sparkutils/testing/TestUtilsEnvironment.scala
@@ -7,6 +7,26 @@ import java.io.File
 
 object TestUtilsEnvironment {
 
+  /**
+   * first attempts to get the system env, then system java property then sqlconf
+   * @param name
+   * @return
+   */
+  def getConfig(name: String, default: String = "") = try {
+    val res = System.getenv(name)
+    if (res ne null)
+      res
+    else {
+      val sp = System.getProperty(name)
+      if (sp ne null)
+        sp
+      else
+        SQLConf.get.getConfString(name, default)
+    }
+  } catch {
+    case _: Throwable => default
+  }
+
   /**
    * Checks if the master is local, if so it's likely tests, if spark.master is not defined it's false
    */
diff --git a/runtime/src/main/scala/com/sparkutils/testing/Testing.scala b/runtime/src/main/scala/com/sparkutils/testing/Testing.scala
@@ -41,4 +41,26 @@ object SparkVersions {
   lazy val sparkVersion = sparkFullVersion.split('.').take(2).mkString(".")
 
   lazy val sparkMajorVersion = sparkFullVersion.split('.').head
+}
+
+
+object ConnectWhenForced {
+
+  val FORCED_CONNECT_PROPERTY_NAME = "SPARKUTILS_TESTING_FORCE_CONNECT"
+
+  /**
+   * When SPARKUTILS_TESTING_FORCE_CONNECT is set to true in ENV or Sys properties then None is returned otherwise thunk.
+   *
+   * This is set/reset by the TestUtils.defaultAndforceConnect helper function
+   *
+   * @param thunk
+   * @tparam T
+   * @return
+   */
+  def someOrForcedConnect[T](thunk: => T): Option[T] =
+    if (TestUtilsEnvironment.getConfig(FORCED_CONNECT_PROPERTY_NAME, "false").toBoolean)
+      None
+    else
+      Some(thunk)
+
 }
diff --git a/testing/src/main/scala/com/sparkutils/testing/TestUtils.scala b/testing/src/main/scala/com/sparkutils/testing/TestUtils.scala
@@ -131,6 +131,26 @@ trait TestUtils extends SessionStrategy with Serializable with ClassicTestUtils
   def not_Cluster(thunk: => Unit): Unit =
     if (TestUtilsEnvironment.shouldRunClusterTests) thunk
 
+  /**
+   * When inConnect is true the thunk is run once with SPARKUTILS_TESTING_FORCE_CONNECT=false and then
+   * with SPARKUTILS_TESTING_FORCE_CONNECT=true enabling functions using someOrForcedConnect to be tested.
+   * @param thunk
+   */
+  def defaultAndforceConnect(thunk: => Unit): Unit =
+    if (inConnect.get()) {
+      val prev = System.getProperty(ConnectWhenForced.FORCED_CONNECT_PROPERTY_NAME) // only possible to change property
+      try {
+        System.setProperty(ConnectWhenForced.FORCED_CONNECT_PROPERTY_NAME, "false")
+        thunk
+        System.setProperty(ConnectWhenForced.FORCED_CONNECT_PROPERTY_NAME, "true")
+        thunk
+      } finally {
+        System.setProperty(ConnectWhenForced.FORCED_CONNECT_PROPERTY_NAME, prev)
+      }
+    } else {
+      thunk // using whatever is current
+    }
+
 }
 
 object TestUtils {
