update ApiPartialEligibilityArm

1 files changed, 11 insertions, 3 deletions

diff --git a/src/main/scala/xyz/driver/pdsuidomain/formats/json/eligibilityarm/ApiPartialEligibilityArm.scala b/src/main/scala/xyz/driver/pdsuidomain/formats/json/eligibilityarm/ApiPartialEligibilityArm.scala
index cc3d460..2041618 100644
--- a/src/main/scala/xyz/driver/pdsuidomain/formats/json/eligibilityarm/ApiPartialEligibilityArm.scala
+++ b/src/main/scala/xyz/driver/pdsuidomain/formats/json/eligibilityarm/ApiPartialEligibilityArm.scala
@@ -1,11 +1,19 @@
 package xyz.driver.pdsuidomain.formats.json.eligibilityarm
 
-import xyz.driver.pdsuidomain.entities.EligibilityArm
+import xyz.driver.pdsuidomain.entities.{EligibilityArm, EligibilityArmDisease, EligibilityArmWithDiseases, Trial}
 import play.api.libs.json.{Format, Json}
 
-final case class ApiPartialEligibilityArm(name: String) {
+final case class ApiPartialEligibilityArm(name: String, diseases: Seq[String]) {
 
-  def applyTo(arm: EligibilityArm): EligibilityArm = arm.copy(name = name)
+  def applyTo(armWithDisease: EligibilityArmWithDiseases): EligibilityArmWithDiseases = {
+    val arm = armWithDisease.eligibilityArm.copy(name = name)
+    val armDiseases = diseases.map { disease =>
+      EligibilityArmDisease(
+        armWithDisease.eligibilityArm.id,
+        Trial.Condition.fromString(disease).getOrElse(throw new NoSuchElementException(s"unknown condition $disease")))
+    }
+    EligibilityArmWithDiseases(arm ,armDiseases)
+  }
 }
 
 object ApiPartialEligibilityArm {