#60: addAdditionalData implementation and test

AbsaOSS · Oct 27, 2023 · 85db371 · 85db371
1 parent 6301732
commit 85db371
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 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
@@ -32,6 +32,7 @@ import scala.collection.immutable.ListMap
  */
 class AtumContext private[agent] (
   val atumPartitions: AtumPartitions,
+  val additionalData: Map[String, Option[String]],
   val agent: AtumAgent,
   private var measures: Set[Measure] = Set.empty
 ) {
@@ -62,8 +63,8 @@ class AtumContext private[agent] (
     )
   }
 
-  def addAdditionalData(key: String, value: String) = {
-    ??? // TODO #60
+  def addAdditionalData(key: String, value: String): AtumContext = {
+    this.copy(additionalData = this.additionalData + (key -> Some(value)))
   }
 
   def addMeasure(newMeasure: Measure): AtumContext = {
@@ -83,10 +84,11 @@ class AtumContext private[agent] (
 
   private[agent] def copy(
     atumPartitions: AtumPartitions = this.atumPartitions,
+    additionalData: Map[String, Option[String]],
     agent: AtumAgent = this.agent,
     measures: Set[Measure] = this.measures
   ): AtumContext = {
-    new AtumContext(atumPartitions, agent, measures)
+    new AtumContext(atumPartitions, additionalData, agent, measures)
   }
 }
 
@@ -114,6 +116,7 @@ object AtumContext {
   private[agent] def fromDTO(atumContextDTO: AtumContextDTO, agent: AtumAgent): AtumContext = {
     new AtumContext(
       AtumPartitions.fromPartitioning(atumContextDTO.partitioning),
+      atumContextDTO.additionalData.additionalData,
       agent,
       MeasuresMapper.mapToMeasures(atumContextDTO.measures)
     )

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
@@ -85,4 +85,18 @@ class AtumContextTest extends AnyFlatSpec with Matchers {
     assert(checkpoint.measurements == measurements)
   }
 
+  "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))
+
+    val newAtumContext = atumContext.addAdditionalData(additionalDataKey, additionalDataValue)
+
+    assert(newAtumContext.additionalData == expectedAdditionalData)
+  }
+
 }