Add for comprehension support to UsingUtils with resource management and guards, and extend test coverage.

cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/utils/UsingUtils.scala

@@ -64,4 +64,63 @@ object UsingUtils {
         }
       }
   }
+
+  /**
+    * Implicits implementing resource management via for comprehension. Example usage:
+    * {{{
+    *   import za.co.absa.cobrix.cobol.utils.UsingUtils.Implicits._
+    *
+    *   for {
+    *     res1 <- new AutoCloseableResource()
+    *     res2 <- new AutoCloseableResource()
+    *   } {
+    *     // Perform operations with the resource
+    *   }
+    * }}}
+    *
+    * or
+    *
+    * {{{
+    *   import za.co.absa.cobrix.cobol.utils.UsingUtils.Implicits._
+    *
+    *   val result = for {
+    *     res1 <- new AutoCloseableResource()
+    *     res2 <- new AutoCloseableResource()
+    *   } yield {
+    *     // Perform operations with the resource, and return a value
+    *   }
+    * }}}
+    */
+  object Implicits {
+    implicit class ResourceWrapper[T <: AutoCloseable](private val resource: T) {
+      def foreach(f: T => Unit): Unit = using(resource)(f)
+
+      def map[U](body: T => U): U = using(resource)(body)
+
+      def flatMap[U](body: T => U): U = using(resource)(body)
+
+      def withFilter(p: T => Boolean): FilteredResourceWrapper[T] = new FilteredResourceWrapper[T](resource, p)
+    }
+  }
+
+  final class FilteredResourceWrapper[T <: AutoCloseable](private val resource: T,
+                                                          private val p: T => Boolean) {
+    def foreach(f: T => Unit): Unit =
+      using(resource) { r =>
+        if (p(r)) f(r) else ()
+      }
+
+    def map[U](body: T => U): U =
+      using(resource) { r =>
+        if (p(r)) body(r) else throw new NoSuchElementException("withFilter predicate is false")
+      }
+
+    def flatMap[U](body: T => U): U =
+      using(resource) { r =>
+        if (p(r)) body(r) else throw new NoSuchElementException("withFilter predicate is false")
+      }
+
+    def withFilter(p2: T => Boolean): FilteredResourceWrapper[T] =
+      new FilteredResourceWrapper[T](resource, r => p(r) && p2(r))
+  }
 }

cobol-parser/src/test/scala/za/co/absa/cobrix/cobol/utils/UsingUtilsSuite.scala

@@ -19,11 +19,13 @@ package za.co.absa.cobrix.cobol.utils
 import org.scalatest.wordspec.AnyWordSpec
 
 class UsingUtilsSuite extends AnyWordSpec {
+  import UsingUtils.Implicits._
+
   "using with a single resource" should {
     "properly close the resource" in {
       var resource: AutoCloseableSpy = null
 
-      UsingUtils.using(new AutoCloseableSpy()) { res =>
+      for (res <- new AutoCloseableSpy()) {
         resource = res
         res.dummyAction()
       }
@@ -179,6 +181,28 @@ class UsingUtilsSuite extends AnyWordSpec {
       assert(resource2.closeCallCount == 1)
     }
 
+    "work with for comprehension" in {
+      var resource1: AutoCloseableSpy = null
+      var resource2: AutoCloseableSpy = null
+
+      val result = for {
+        res1 <- new AutoCloseableSpy()
+        res2 <- new AutoCloseableSpy()
+      } yield {
+        resource1 = res1
+        resource2 = res2
+        res1.dummyAction()
+        res2.dummyAction()
+        100
+      }
+
+      assert(result == 100)
+      assert(resource1.actionCallCount == 1)
+      assert(resource1.closeCallCount == 1)
+      assert(resource2.actionCallCount == 1)
+      assert(resource2.closeCallCount == 1)
+    }
+
     "properly close both resources when an inner one throws an exception during action and close" in {
       var resource1: AutoCloseableSpy = null
       var resource2: AutoCloseableSpy = null
@@ -209,6 +233,37 @@ class UsingUtilsSuite extends AnyWordSpec {
       assert(resource2.closeCallCount == 1)
     }
 
+    "properly close both resources when an inner one throws an exception during action and close (for comprehension)" in {
+      var resource1: AutoCloseableSpy = null
+      var resource2: AutoCloseableSpy = null
+      var exceptionThrown = false
+
+      try {
+        for {
+          res1 <- new AutoCloseableSpy()
+          res2 <- new AutoCloseableSpy(failAction = true, failClose = true)
+        } {
+          resource1 = res1
+          resource2 = res2
+          res1.dummyAction()
+          res2.dummyAction()
+        }
+      } catch {
+        case ex: Throwable =>
+          exceptionThrown = true
+          assert(ex.getMessage.contains("Failed during action"))
+          val suppressed = ex.getSuppressed
+          assert(suppressed.length == 1)
+          assert(suppressed(0).getMessage.contains("Failed to close resource"))
+      }
+
+      assert(exceptionThrown)
+      assert(resource1.actionCallCount == 1)
+      assert(resource1.closeCallCount == 1)
+      assert(resource2.actionCallCount == 1)
+      assert(resource2.closeCallCount == 1)
+    }
+
     "properly close both resources when an outer one throws an exception during action and close" in {
       var resource1: AutoCloseableSpy = null
       var resource2: AutoCloseableSpy = null
@@ -239,6 +294,37 @@ class UsingUtilsSuite extends AnyWordSpec {
       assert(resource2.closeCallCount == 1)
     }
 
+    "properly close both resources when an outer one throws an exception during action and close (for comprehension)" in {
+      var resource1: AutoCloseableSpy = null
+      var resource2: AutoCloseableSpy = null
+      var exceptionThrown = false
+
+      try {
+        for {
+          res1 <- new AutoCloseableSpy(failAction = true, failClose = true)
+          res2 <- new AutoCloseableSpy()
+        } {
+          resource1 = res1
+          resource2 = res2
+          res1.dummyAction()
+          res2.dummyAction()
+        }
+      } catch {
+        case ex: Throwable =>
+          exceptionThrown = true
+          assert(ex.getMessage.contains("Failed during action"))
+          val suppressed = ex.getSuppressed
+          assert(suppressed.length == 1)
+          assert(suppressed(0).getMessage.contains("Failed to close resource"))
+      }
+
+      assert(exceptionThrown)
+      assert(resource1.actionCallCount == 1)
+      assert(resource1.closeCallCount == 1)
+      assert(resource2.actionCallCount == 0)
+      assert(resource2.closeCallCount == 1)
+    }
+
     "properly close the outer resource when the inner one fails on create" in {
       var resource1: AutoCloseableSpy = null
       var resource2: AutoCloseableSpy = null
@@ -264,5 +350,163 @@ class UsingUtilsSuite extends AnyWordSpec {
       assert(resource1.closeCallCount == 1)
       assert(resource2 == null)
     }
+
+    "properly close the outer resource when the inner one fails on create (for comprehension)" in {
+      val resource1: AutoCloseableSpy = new AutoCloseableSpy()
+      var resource2: AutoCloseableSpy = null
+      var exceptionThrown = false
+
+      try {
+        resource1
+          .flatMap(res1 =>
+            new AutoCloseableSpy(failCreate = true)
+              .map(res2 => {
+                resource2 = res2
+                res1.dummyAction()
+                res2.dummyAction()
+              })
+          )
+      } catch {
+        case ex: Throwable =>
+          exceptionThrown = true
+          assert(ex.getMessage.contains("Failed to create resource"))
+      }
+
+      assert(exceptionThrown)
+      assert(resource1.actionCallCount == 0)
+      assert(resource1.closeCallCount == 1)
+      assert(resource2 == null)
+    }
+  }
+
+  "withFilter (for-comprehension guards)" should {
+    "skip the body when the guard is false (foreach form) and still close the resource" in {
+      val res = new AutoCloseableSpy()
+
+      var bodyRan = false
+      for {
+        r <- res if false
+      } {
+        bodyRan = true
+        r.dummyAction()
+      }
+
+      assert(!bodyRan)
+      assert(res.actionCallCount == 0)
+      assert(res.closeCallCount == 1)
+    }
+
+    "run the body when the guard is true (foreach form) and close the resource" in {
+      val res = new AutoCloseableSpy()
+
+      var bodyRan = false
+      for {
+        r <- res if true
+      } {
+        bodyRan = true
+        r.dummyAction()
+      }
+
+      assert(bodyRan)
+      assert(res.actionCallCount == 1)
+      assert(res.closeCallCount == 1)
+    }
+
+    "close both resources when an inner guard is false (nested generators)" in {
+      val r1 = new AutoCloseableSpy()
+      val r2 = new AutoCloseableSpy()
+
+      var bodyRan = false
+      for {
+        a <- r1
+        b <- r2 if false
+      } {
+        bodyRan = true
+        a.dummyAction()
+        b.dummyAction()
+      }
+
+      assert(!bodyRan)
+      assert(r1.actionCallCount == 0)
+      assert(r2.actionCallCount == 0)
+      assert(r2.closeCallCount == 1)
+      assert(r1.closeCallCount == 1)
+    }
+
+    "compose multiple guards correctly (both must be true)" in {
+      val res = new AutoCloseableSpy()
+
+      var bodyRan = false
+      for {
+        r <- res if true if false
+      } {
+        bodyRan = true
+        r.dummyAction()
+      }
+
+      assert(!bodyRan)
+      assert(res.actionCallCount == 0)
+      assert(res.closeCallCount == 1)
+    }
+
+    "not evaluate the body when the first guard is false (side-effect check)" in {
+      val res = new AutoCloseableSpy()
+
+      var sideEffect = 0
+      for {
+        r <- res if false if { sideEffect += 1; true }
+      } {
+        r.dummyAction()
+      }
+
+      assert(sideEffect == 0)
+      assert(res.actionCallCount == 0)
+      assert(res.closeCallCount == 1)
+    }
+
+    "throw on yield when the guard is false (map path) and still close the resource" in {
+      val res = new AutoCloseableSpy()
+
+      var thrown = false
+      try {
+        val _ = for {
+          r <- res if false
+        } yield {
+          r.dummyAction()
+          1
+        }
+      } catch {
+        case _: NoSuchElementException => thrown = true
+      }
+
+      assert(thrown)
+      assert(res.actionCallCount == 0)
+      assert(res.closeCallCount == 1)
+    }
+
+    "throw on a guarded middle generator in a yield (flatMap->map path) and close all opened resources" in {
+      val r1 = new AutoCloseableSpy()
+      val r2 = new AutoCloseableSpy()
+
+      var thrown = false
+      try {
+        val _ = for {
+          a <- r1
+          b <- r2 if false
+        } yield {
+          a.dummyAction()
+          b.dummyAction()
+          1
+        }
+      } catch {
+        case _: NoSuchElementException => thrown = true
+      }
+
+      assert(thrown)
+      assert(r1.actionCallCount == 0)
+      assert(r2.actionCallCount == 0)
+      assert(r2.closeCallCount == 1)
+      assert(r1.closeCallCount == 1)
+    }
   }
 }