From 067b15c34e4aa4f7756c42be17e3b77e932490b5 Mon Sep 17 00:00:00 2001
From: Daniel Slapman <danslapman@gmail.com>
Date: Sun, 10 Mar 2024 22:57:53 +0100
Subject: [PATCH] Support templates in response headers

---
 .../tcb/mockingbird/api/PublicApiHandler.scala      | 11 +++++++----
 .../tinkoff/tcb/utils/transformation/package.scala  |  4 ++--
 .../tcb/utils/transformation/string/package.scala   |  3 +++
 .../string/StringTransformationsSpec.scala          | 13 +++++++++++++
 .../scala/ru/tinkoff/tcb/utils/any/package.scala    |  8 ++++++++
 5 files changed, 33 insertions(+), 6 deletions(-)
 create mode 100644 backend/utils/src/main/scala/ru/tinkoff/tcb/utils/any/package.scala

diff --git a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/api/PublicApiHandler.scala b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/api/PublicApiHandler.scala
index 024f96f6..b5826b8b 100644
--- a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/api/PublicApiHandler.scala
+++ b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/mockingbird/api/PublicApiHandler.scala
@@ -43,6 +43,7 @@ import ru.tinkoff.tcb.mockingbird.model.XmlProxyResponse
 import ru.tinkoff.tcb.mockingbird.scenario.CallbackEngine
 import ru.tinkoff.tcb.mockingbird.scenario.ScenarioEngine
 import ru.tinkoff.tcb.protocol.log.*
+import ru.tinkoff.tcb.utils.any.*
 import ru.tinkoff.tcb.utils.circe.optics.JsonOptic
 import ru.tinkoff.tcb.utils.regex.*
 import ru.tinkoff.tcb.utils.sandboxing.GraalJsSandbox
@@ -128,11 +129,13 @@ final class PublicApiHandler(
           )
         case _ =>
           ZIO.succeed(
-            if (stub.response.isTemplate) {
-              HttpStubResponse.jsonBody
+            stub.response
+              .applyIf(_.isTemplate)(HttpStubResponse.jsonBody
                 .updateF(_.substitute(data).substitute(xdata))
-                .andThen(HttpStubResponse.xmlBody.updateF(_.substitute(data).substitute(xdata)))(stub.response)
-            } else stub.response
+                .andThen(HttpStubResponse.xmlBody.updateF(_.substitute(data).substitute(xdata)))
+              )
+              .applyIf(HttpStubResponse.headers.getOption(_).exists(_.values.exists(_.isTemplate)))
+                (HttpStubResponse.headers.updateF(_.view.mapValues(_.substitute(data, xdata)).toMap))
           )
       }
       _ <- ZIO.when(stub.scope == Scope.Countdown)(stubDAO.updateById(stub.id, prop[HttpStub](_.times).inc(-1)))
diff --git a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/package.scala b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/package.scala
index 7b28f829..0d2b5d25 100644
--- a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/package.scala
+++ b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/package.scala
@@ -3,6 +3,6 @@ package ru.tinkoff.tcb.utils
 import scala.util.matching.Regex
 
 package object transformation {
-  val SubstRx: Regex = """\$\{(.*?)\}""".r
-  val CodeRx: Regex  = """%\{(.*?)\}""".r
+  val SubstRx: Regex = """\$\{(.+?)\}""".r
+  val CodeRx: Regex  = """%\{(.+?)\}""".r
 }
diff --git a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/string/package.scala b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/string/package.scala
index 199e5298..ba08f992 100644
--- a/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/string/package.scala
+++ b/backend/mockingbird/src/main/scala/ru/tinkoff/tcb/utils/transformation/string/package.scala
@@ -8,6 +8,9 @@ import ru.tinkoff.tcb.utils.transformation.json.*
 
 package object string {
   implicit final class StringTransformations(private val s: String) extends AnyVal {
+    def isTemplate: Boolean =
+      CodeRx.findFirstIn(s).isDefined || SubstRx.findFirstIn(s).isDefined
+
     def substitute(jvalues: Json, xvalues: Node)(implicit sandbox: GraalJsSandbox): String =
       if (SubstRx.findFirstIn(s).isDefined || CodeRx.findFirstIn(s).isDefined)
         Json.fromString(s).substitute(jvalues).substitute(xvalues).asString.getOrElse(s)
diff --git a/backend/mockingbird/src/test/scala/ru/tinkoff/tcb/utils/transformation/string/StringTransformationsSpec.scala b/backend/mockingbird/src/test/scala/ru/tinkoff/tcb/utils/transformation/string/StringTransformationsSpec.scala
index e0a9ffde..d63fd781 100644
--- a/backend/mockingbird/src/test/scala/ru/tinkoff/tcb/utils/transformation/string/StringTransformationsSpec.scala
+++ b/backend/mockingbird/src/test/scala/ru/tinkoff/tcb/utils/transformation/string/StringTransformationsSpec.scala
@@ -24,4 +24,17 @@ class StringTransformationsSpec extends AnyFunSuite with Matchers with TryValues
 
     "${/a}".substitute(Json.Null, xml("<a>test</a>")) shouldBe "test"
   }
+
+  test("isTemplate test") {
+    "".isTemplate shouldBe false
+    "{}".isTemplate shouldBe false
+    "${}".isTemplate shouldBe false
+    "${a}".isTemplate shouldBe true
+    "${a.b}".isTemplate shouldBe true
+    "${a.[0]}".isTemplate shouldBe true
+    "${a.[0].b}".isTemplate shouldBe true
+
+    "%{}".isTemplate shouldBe false
+    "%{var a = 1}".isTemplate shouldBe true
+  }
 }
diff --git a/backend/utils/src/main/scala/ru/tinkoff/tcb/utils/any/package.scala b/backend/utils/src/main/scala/ru/tinkoff/tcb/utils/any/package.scala
new file mode 100644
index 00000000..584d0df0
--- /dev/null
+++ b/backend/utils/src/main/scala/ru/tinkoff/tcb/utils/any/package.scala
@@ -0,0 +1,8 @@
+package ru.tinkoff.tcb.utils
+
+package object any {
+  implicit class AnyExtensionOps[T](private val t: T) extends AnyVal {
+    @inline def applyIf(condition: T => Boolean)(fun: T => T): T =
+      if (condition(t)) fun(t) else t
+  }
+}