Merge remote-tracking branch 'origin/master' into feature/26-checkpoi…

…nt-creation-function # Conflicts: # agent/src/main/scala/za/co/absa/atum/agent/AtumContext.scala
AbsaOSS · Oct 30, 2023 · a31fd04 · a31fd04
2 parents c2c3405 + 7681d9d
commit a31fd04
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 6 deletions.
diff --git a/agent/src/main/scala/za/co/absa/atum/agent/AtumContext.scala b/agent/src/main/scala/za/co/absa/atum/agent/AtumContext.scala
@@ -34,7 +34,8 @@ import scala.collection.immutable.ListMap
 class AtumContext private[agent] (
   val atumPartitions: AtumPartitions,
   val agent: AtumAgent,
-  private var measures: Set[Measure] = Set.empty
+  private var measures: Set[Measure] = Set.empty,
+  private var additionalData: Map[String, Option[String]] = Map.empty
 ) {
 
   def currentMeasures: Set[Measure] = measures
@@ -87,8 +88,12 @@ class AtumContext private[agent] (
     this
   }
 
-  def addAdditionalData(key: String, value: String) = {
-    ??? // TODO #60
+  def addAdditionalData(key: String, value: String): Unit = {
+    additionalData += (key -> Some(value))
+  }
+
+  def currentAdditionalData: Map[String, Option[String]] = {
+    this.additionalData
   }
 
   def addMeasure(newMeasure: Measure): AtumContext = {
@@ -109,9 +114,10 @@ class AtumContext private[agent] (
   private[agent] def copy(
     atumPartitions: AtumPartitions = this.atumPartitions,
     agent: AtumAgent = this.agent,
-    measures: Set[Measure] = this.measures
+    measures: Set[Measure] = this.measures,
+    additionalData: Map[String, Option[String]] = this.additionalData
   ): AtumContext = {
-    new AtumContext(atumPartitions, agent, measures)
+    new AtumContext(atumPartitions, agent, measures, additionalData)
   }
 }
 
@@ -140,7 +146,8 @@ object AtumContext {
     new AtumContext(
       AtumPartitions.fromPartitioning(atumContextDTO.partitioning),
       agent,
-      MeasuresBuilder.mapToMeasures(atumContextDTO.measures)
+      MeasuresBuilder.mapToMeasures(atumContextDTO.measures),
+      atumContextDTO.additionalData.additionalData
     )
   }
 

diff --git a/agent/src/test/scala/za/co/absa/atum/agent/AtumContextTest.scala b/agent/src/test/scala/za/co/absa/atum/agent/AtumContextTest.scala
@@ -178,4 +178,18 @@ class AtumContextTest extends AnyFlatSpec with Matchers {
     assert(argumentSecond.getValue.measurements.tail.head.result.mainValue.valueType == ResultValueType.BigDecimal)
   }
 
+  "addAdditionalData" should "add key/value pair to map for additional data" in {
+    val atumAgent = new AtumAgent
+    val atumPartitions = AtumPartitions("key" -> "val")
+    val atumContext = atumAgent.getOrCreateAtumContext(atumPartitions)
+
+    val additionalDataKey = "additionalKey"
+    val additionalDataValue = "additionalVal"
+    val expectedAdditionalData =  Map(additionalDataKey -> Some(additionalDataValue))
+
+    atumContext.addAdditionalData(additionalDataKey, additionalDataValue)
+
+    assert(atumContext.currentAdditionalData == expectedAdditionalData)
+  }
+
 }