1 files changed, 23 insertions, 16 deletions
diff --git a/src/main/scala/xyz/driver/pdsuidomain/services/fake/FakeTrialService.scala b/src/main/scala/xyz/driver/pdsuidomain/services/fake/FakeTrialService.scala
index 9253f88..0431d95 100644
--- a/src/main/scala/xyz/driver/pdsuidomain/services/fake/FakeTrialService.scala
+++ b/src/main/scala/xyz/driver/pdsuidomain/services/fake/FakeTrialService.scala
@@ -7,11 +7,12 @@ import akka.stream.scaladsl.Source
 import akka.util.ByteString
 import xyz.driver.core.generators
 import xyz.driver.entities.labels.Label
+import xyz.driver.entities.patient.CancerType
 import xyz.driver.pdsuicommon.auth.AuthenticatedRequestContext
 import xyz.driver.pdsuicommon.db._
 import xyz.driver.pdsuicommon.domain.{LongId, StringId, UuidId}
-import xyz.driver.pdsuidomain.entities.{Arm, Criterion, Trial}
 import xyz.driver.pdsuidomain.entities.export.trial.{ExportTrialArm, ExportTrialLabelCriterion, ExportTrialWithLabels}
+import xyz.driver.pdsuidomain.entities.{Criterion, EligibilityArm, Trial}
 import xyz.driver.pdsuidomain.services.TrialService
 
 import scala.concurrent.Future
@@ -56,30 +57,36 @@ class FakeTrialService extends TrialService {
           implicit requestContext: AuthenticatedRequestContext): Future[GetListReply] =
     Future.successful(GetListReply.EntityList(Seq(trial), 1, None))
 
-  override def getTrialWithLabels(trialId: StringId[Trial], condition: String)(
-          implicit requestContext: AuthenticatedRequestContext): Future[GetTrialWithLabelsReply] = {
-    Future.successful(
-      GetTrialWithLabelsReply.Entity(ExportTrialWithLabels(
-        StringId[Trial]("NCT" + generators.nextInt(999999).toString),
-        UuidId[Trial](generators.nextUuid()),
-        generators.oneOf("adenocarcinoma", "breast", "prostate"),
-        LocalDateTime.now(),
-        labelVersion = 1L,
-        generators.listOf(new ExportTrialArm(
-          LongId[Arm](generators.nextInt(999999).toLong),
+  override def getTrialWithLabels(trialId: StringId[Trial], cancerType: CancerType)(
+          implicit requestContext: AuthenticatedRequestContext): Future[GetTrialWithLabelsReply] =
+    Future.successful(GetTrialWithLabelsReply.Entity(nextExportTrialWithLabels()))
+
+  override def getTrialsWithLabels(cancerType: CancerType)(
+          implicit requestContext: AuthenticatedRequestContext): Future[GetTrialsWithLabelsReply] =
+    Future.successful(GetTrialsWithLabelsReply.EntityList(generators.seqOf(nextExportTrialWithLabels())))
+
+  private def nextExportTrialWithLabels() =
+    ExportTrialWithLabels(
+      StringId[Trial]("NCT" + generators.nextInt(999999).toString),
+      UuidId[Trial](generators.nextUuid()),
+      LocalDateTime.now(),
+      labelVersion = 1L,
+      generators.listOf(
+        new ExportTrialArm(
+          LongId[EligibilityArm](generators.nextInt(999999).toLong),
           generators.nextName().value
         )),
-        generators.listOf(new ExportTrialLabelCriterion(
+      generators.listOf(
+        new ExportTrialLabelCriterion(
           LongId[Criterion](generators.nextInt(999999).toLong),
           generators.nextOption(generators.nextBoolean()),
           LongId[Label](generators.nextInt(999999).toLong),
-          generators.setOf(LongId[Arm](generators.nextInt(999999).toLong)),
+          generators.setOf(LongId[EligibilityArm](generators.nextInt(999999).toLong)),
           generators.nextName().value,
           generators.nextBoolean(),
           generators.nextBoolean()
         ))
-      )))
-  }
+    )
 
   def update(origTrial: Trial, draftTrial: Trial)(
           implicit requestContext: AuthenticatedRequestContext): Future[UpdateReply] =