refactor commonality in RelationsServiceTest (#2791)

wellcomecollection · Dec 19, 2024 · e1dd763 · e1dd763
1 parent 2622c78
commit e1dd763
Showing 1 changed file with 125 additions and 138 deletions.
diff --git a/...lation_embedder/src/test/scala/weco/pipeline/relation_embedder/RelationsServiceTest.scala b/...lation_embedder/src/test/scala/weco/pipeline/relation_embedder/RelationsServiceTest.scala
@@ -10,6 +10,7 @@ import weco.catalogue.internal_model.Implicits._
 import weco.catalogue.internal_model.fixtures.index.IndexFixtures
 import weco.catalogue.internal_model.work.WorkState.Merged
 import weco.catalogue.internal_model.work.Work
+import weco.fixtures.TestWith
 import weco.pipeline.relation_embedder.fixtures.RelationGenerators
 import weco.pipeline.relation_embedder.models._
 
@@ -68,72 +69,88 @@ class RelationsServiceTest
       work3,
       work4
     )
+  protected def withLocalMergedWorksIndexContaining[R](
+    indexedWorks: Seq[Work[Merged]]
+  )(
+    testWith: TestWith[
+      Index,
+      R
+    ]
+  ): R = {
+    withLocalMergedWorksIndex {
+      index =>
+        insertIntoElasticsearch(index, indexedWorks: _*)
+        testWith(index)
+    }
+  }
 
   describe("getAffectedWorks") {
 
     import Selector._
 
+    def queryAffectedWorks(
+      index: Index,
+      batch: Batch,
+      completeTreeScroll: Int = 20,
+      affectedWorksScroll: Int = 20
+    ): Seq[Work[Merged]] = withActorSystem {
+      implicit actorSystem =>
+        whenReady(
+          service(
+            index,
+            completeTreeScroll = completeTreeScroll,
+            affectedWorksScroll = affectedWorksScroll
+          ).getAffectedWorks(batch).runWith(Sink.seq[Work[Merged]])
+        ) {
+          foundWorks: Seq[Work[Merged]] => foundWorks
+        }
+    }
+
     it("Retrieves all affected works when batch consists of a complete tree") {
-      withLocalMergedWorksIndex {
+      withLocalMergedWorksIndexContaining(works) {
         index =>
-          insertIntoElasticsearch(index, works: _*)
-          withActorSystem {
-            implicit actorSystem =>
-              val batch = Batch(rootPath = "A", List(Tree("A")))
-              whenReady(queryAffectedWorks(service(index), batch)) {
-                _ should contain theSameElementsAs works
-              }
-          }
+          queryAffectedWorks(
+            index,
+            Batch(rootPath = "A", List(Tree("A")))
+          ) should contain theSameElementsAs works
       }
     }
 
     it("Retrieves all affected works when batch consists of single node") {
-      withLocalMergedWorksIndex {
+      withLocalMergedWorksIndexContaining(works) {
         index =>
-          insertIntoElasticsearch(index, works: _*)
-          withActorSystem {
-            implicit actorSystem =>
-              val batch = Batch(rootPath = "A", List(Node("A/B")))
-              whenReady(queryAffectedWorks(service(index), batch)) {
-                _ should contain theSameElementsAs List(workB)
-              }
-          }
+          queryAffectedWorks(
+            index,
+            Batch(rootPath = "A", List(Node("A/B")))
+          ) should contain theSameElementsAs List(workB)
       }
     }
 
     it("Retrieves all affected works when batch consists of a nodes children") {
-      withLocalMergedWorksIndex {
+      withLocalMergedWorksIndexContaining(works) {
         index =>
-          insertIntoElasticsearch(index, works: _*)
-          withActorSystem {
-            implicit actorSystem =>
-              val batch = Batch(rootPath = "A", List(Children("A/C")))
-              whenReady(queryAffectedWorks(service(index), batch)) {
-                _ should contain theSameElementsAs List(workX, workY, workZ)
-              }
-          }
+          queryAffectedWorks(
+            index,
+            Batch(rootPath = "A", List(Children("A/C")))
+          ) should contain theSameElementsAs List(workX, workY, workZ)
       }
     }
 
     it(
       "Retrieves all affected works when batch consists of a nodes descendents"
     ) {
-      withLocalMergedWorksIndex {
+      withLocalMergedWorksIndexContaining(works) {
         index =>
-          insertIntoElasticsearch(index, works: _*)
-          withActorSystem {
-            implicit actorSystem =>
-              val batch = Batch(rootPath = "A", List(Descendents("A/C")))
-              whenReady(queryAffectedWorks(service(index), batch)) {
-                _ should contain theSameElementsAs List(
-                  workX,
-                  workY,
-                  workZ,
-                  work3,
-                  work4
-                )
-              }
-          }
+          queryAffectedWorks(
+            index,
+            Batch(rootPath = "A", List(Descendents("A/C")))
+          ) should contain theSameElementsAs List(
+            workX,
+            workY,
+            workZ,
+            work3,
+            work4
+          )
       }
     }
 
@@ -143,75 +160,56 @@ class RelationsServiceTest
       withLocalMergedWorksIndex {
         index =>
           insertIntoElasticsearch(index, works: _*)
-          withActorSystem {
-            implicit actorSystem =>
-              val batch = Batch(
-                rootPath = "A",
-                List(Node("A/E/2"), Descendents("A/C"), Children("A"))
-              )
-              whenReady(queryAffectedWorks(service(index), batch)) {
-                _ should contain theSameElementsAs List(
-                  workB,
-                  workC,
-                  workE,
-                  workX,
-                  workY,
-                  workZ,
-                  work2,
-                  work3,
-                  work4
-                )
-              }
-          }
+          val batch = Batch(
+            rootPath = "A",
+            List(Node("A/E/2"), Descendents("A/C"), Children("A"))
+          )
+          queryAffectedWorks(
+            index,
+            batch
+          ) should contain theSameElementsAs List(
+            workB,
+            workC,
+            workE,
+            workX,
+            workY,
+            workZ,
+            work2,
+            work3,
+            work4
+          )
       }
     }
 
     it("Retrieves all affected works across multiple scroll pages") {
-      withLocalMergedWorksIndex {
+      withLocalMergedWorksIndexContaining(works) {
         index =>
-          insertIntoElasticsearch(index, works: _*)
-          withActorSystem {
-            implicit actorSystem =>
-              val batch = Batch(rootPath = "A", List(Tree("A")))
-              whenReady(
-                queryAffectedWorks(
-                  service(index, affectedWorksScroll = 3),
-                  batch
-                )
-              ) {
-                _ should contain theSameElementsAs works
-              }
-          }
+          queryAffectedWorks(
+            index,
+            Batch(rootPath = "A", List(Tree("A"))),
+            affectedWorksScroll = 3
+          ) should contain theSameElementsAs works
       }
     }
 
     it("Returns invisible works") {
-      withLocalMergedWorksIndex {
+      val invisibleWork = work("A/C/X/5").invisible()
+      withLocalMergedWorksIndexContaining(invisibleWork :: works) {
         index =>
-          val invisibleWork = work("A/C/X/5").invisible()
-          insertIntoElasticsearch(index, invisibleWork :: works: _*)
-          withActorSystem {
-            implicit actorSystem =>
-              val batch = Batch(
-                rootPath = "A",
-                List(Children("A/C/X"), Descendents("A/C/X"), Node("A/C/X/5"))
-              )
-              whenReady(queryAffectedWorks(service(index), batch)) {
-                _ should contain theSameElementsAs List(
-                  work3,
-                  work4,
-                  invisibleWork
-                )
-              }
-          }
+          val batch = Batch(
+            rootPath = "A",
+            List(Children("A/C/X"), Descendents("A/C/X"), Node("A/C/X/5"))
+          )
+          queryAffectedWorks(
+            index,
+            batch
+          ) should contain theSameElementsAs List(
+            work3,
+            work4,
+            invisibleWork
+          )
       }
     }
-
-    def queryAffectedWorks(service: RelationsService, batch: Batch)(
-      implicit as: ActorSystem
-    ) =
-      service.getAffectedWorks(batch).runWith(Sink.seq[Work[Merged]])
-
   }
 
   describe("getRelationTree") {
@@ -221,44 +219,28 @@ class RelationsServiceTest
     val batch = Batch("A", List(Children("A/B"), Node("A/C/X")))
 
     it("Retrieves all works in archive") {
-      withLocalMergedWorksIndex {
+      withLocalMergedWorksIndexContaining(works) {
         index =>
-          withActorSystem {
-            implicit actorSystem =>
-              insertIntoElasticsearch(index, works: _*)
-              whenReady(queryRelationTree(service(index), batch)) {
-                _ should contain theSameElementsAs works.map(toRelationWork)
-              }
-          }
+          queryRelationTree(index, batch) should contain theSameElementsAs works
+            .map(toRelationWork)
       }
     }
 
     it("Ignores works in other archives") {
-      withLocalMergedWorksIndex {
+      withLocalMergedWorksIndexContaining(work("other/archive") :: works) {
         index =>
-          withActorSystem {
-            implicit actorSystem =>
-              insertIntoElasticsearch(index, work("other/archive") :: works: _*)
-              whenReady(queryRelationTree(service(index), batch)) {
-                _ should contain theSameElementsAs works.map(toRelationWork)
-              }
-          }
+          queryRelationTree(index, batch) should contain theSameElementsAs works
+            .map(toRelationWork)
       }
     }
 
     it("Ignores invisible works") {
-      withLocalMergedWorksIndex {
+      withLocalMergedWorksIndexContaining(
+        work("A/Invisible").invisible() :: works
+      ) {
         index =>
-          withActorSystem {
-            implicit actorSystem =>
-              insertIntoElasticsearch(
-                index,
-                work("A/Invisible").invisible() :: works: _*
-              )
-              whenReady(queryRelationTree(service(index), batch)) {
-                _ should contain theSameElementsAs works.map(toRelationWork)
-              }
-          }
+          queryRelationTree(index, batch) should contain theSameElementsAs works
+            .map(toRelationWork)
       }
     }
 
@@ -276,20 +258,25 @@ class RelationsServiceTest
         selectors = List(Tree("x"))
       )
       val expected = works.map(toRelationWork)
-      withLocalMergedWorksIndex {
+      withLocalMergedWorksIndexContaining(works) {
         index =>
-          withActorSystem {
-            implicit actorSystem =>
-              insertIntoElasticsearch(index, works: _*)
-              whenReady(queryRelationTree(service(index), batch)) {
-                _ should contain theSameElementsAs expected
-              }
-          }
+          queryRelationTree(
+            index,
+            batch
+          ) should contain theSameElementsAs expected
       }
     }
-    def queryRelationTree(service: RelationsService, batch: Batch)(
-      implicit as: ActorSystem
-    ) =
-      service.getRelationTree(batch).runWith(Sink.seq[RelationWork])
+
+    def queryRelationTree(index: Index, batch: Batch): Seq[RelationWork] =
+      withActorSystem {
+        implicit actorSystem: ActorSystem =>
+          whenReady(
+            service(index)
+              .getRelationTree(batch)
+              .runWith(Sink.seq[RelationWork])
+          ) {
+            relationWorks: Seq[RelationWork] => relationWorks
+          }
+      }
   }
 }