1 files changed, 5 insertions, 6 deletions
diff --git a/src/test/scala/xyz/driver/pdsuidomain/formats/json/sprayformats/ExportFormatSuite.scala b/src/test/scala/xyz/driver/pdsuidomain/formats/json/sprayformats/ExportFormatSuite.scala
index 402e893..d78e754 100644
--- a/src/test/scala/xyz/driver/pdsuidomain/formats/json/sprayformats/ExportFormatSuite.scala
+++ b/src/test/scala/xyz/driver/pdsuidomain/formats/json/sprayformats/ExportFormatSuite.scala
@@ -3,8 +3,8 @@ package xyz.driver.pdsuidomain.formats.json.sprayformats
 import java.time.{LocalDate, LocalDateTime}
 import java.util.UUID
 
-import spray.json._
 import org.scalatest.{FlatSpec, Matchers}
+import spray.json._
 import xyz.driver.entities.labels.LabelValue
 import xyz.driver.pdsuicommon.domain.{LongId, StringId, UuidId}
 import xyz.driver.pdsuidomain.entities.{DocumentType, ProviderType, RecordRequestId}
@@ -81,8 +81,8 @@ class ExportFormatSuite extends FlatSpec with Matchers {
   "Json format for ApiExportTrialWithLabels" should "read and write correct JSON" in {
     import xyz.driver.pdsuidomain.entities.export.trial._
     val arms = List(
-      ExportTrialArm(armId = LongId(1), armName = "arm 1"),
-      ExportTrialArm(armId = LongId(2), armName = "arm 2")
+      ExportTrialArm(armId = LongId(1), armName = "arm 1", diseaseList = Seq("Breast")),
+      ExportTrialArm(armId = LongId(2), armName = "arm 2", diseaseList = Seq("Breast"))
     )
     val criteriaList = List(
       ExportTrialLabelCriterion(
@@ -107,7 +107,6 @@ class ExportFormatSuite extends FlatSpec with Matchers {
     val trialWithLabels = ExportTrialWithLabels(
       nctId = StringId("NCT000001"),
       trialId = UuidId("40892a07-c638-49d2-9795-1edfefbbcc7c"),
-      disease = "Breast",
       lastReviewed = LocalDateTime.parse("2017-08-10T18:00:00"),
       labelVersion = 1L,
       arms = arms,
@@ -116,8 +115,8 @@ class ExportFormatSuite extends FlatSpec with Matchers {
 
     val writtenJson = trialWithLabelsFormat.write(trialWithLabels)
     writtenJson should be(
-      """{"nctId":"NCT000001","trialId":"40892a07-c638-49d2-9795-1edfefbbcc7c","disease":"Breast","lastReviewed":"2017-08-10T18:00Z",
-        "labelVersion":1,"arms":[{"armId":1,"armName":"arm 1"},{"armId":2,"armName":"arm 2"}],"criteria":[
+      """{"nctId":"NCT000001","trialId":"40892a07-c638-49d2-9795-1edfefbbcc7c","lastReviewed":"2017-08-10T18:00Z",
+        "labelVersion":1,"arms":[{"armId":1,"armName":"arm 1","diseaseList":["Breast"]},{"armId":2,"armName":"arm 2","diseaseList":["Breast"]}],"criteria":[
         {"value":"Yes","labelId":21,"criterionId":10,"criterionText":"criteria 10 text","armIds":[1,2],"isCompound":false,"isDefining":false},
         {"value":"Unknown","labelId":21,"criterionId":11,"criterionText":"criteria 11 text","armIds":[2],"isCompound":true,"isDefining":false}]}""".parseJson)
   }