From 1bfad8f92eda93ae8196b887851eb2c25ff391d4 Mon Sep 17 00:00:00 2001
From: Andreas Lochbihler <andreas.lochbihler@digitalasset.com>
Date: Thu, 12 Dec 2024 17:15:21 +0100
Subject: [PATCH] basic LAPI TT tests for prefetchContractKeys

---
 .../ledger/api/testtool/suites/v1_17.scala    |  2 +-
 .../suites/v1_17/PrefetchContractKeysIT.scala | 82 +++++++++++++++++++
 2 files changed, 83 insertions(+), 1 deletion(-)
 create mode 100644 sdk/ledger-test-tool/suites/src/main/scala/com/daml/ledger/api/testtool/suites/v1_17/PrefetchContractKeysIT.scala

diff --git a/sdk/ledger-test-tool/suites/src/main/scala/com/daml/ledger/api/testtool/suites/v1_17.scala b/sdk/ledger-test-tool/suites/src/main/scala/com/daml/ledger/api/testtool/suites/v1_17.scala
index a444d4d25c42..fc97c21775e8 100644
--- a/sdk/ledger-test-tool/suites/src/main/scala/com/daml/ledger/api/testtool/suites/v1_17.scala
+++ b/sdk/ledger-test-tool/suites/src/main/scala/com/daml/ledger/api/testtool/suites/v1_17.scala
@@ -8,7 +8,7 @@ import com.daml.ledger.api.tls.TlsConfiguration
 
 package object v1_17 {
   def default(timeoutScaleFactor: Double): Vector[LedgerTestSuite] =
-    v1_15.default(timeoutScaleFactor) ++ Vector(new UpgradingIT)
+    v1_15.default(timeoutScaleFactor) ++ Vector(new UpgradingIT, new PrefetchContractKeysIT)
 
   def optional(tlsConfig: Option[TlsConfiguration]): Vector[LedgerTestSuite] =
     v1_15.optional(tlsConfig)
diff --git a/sdk/ledger-test-tool/suites/src/main/scala/com/daml/ledger/api/testtool/suites/v1_17/PrefetchContractKeysIT.scala b/sdk/ledger-test-tool/suites/src/main/scala/com/daml/ledger/api/testtool/suites/v1_17/PrefetchContractKeysIT.scala
new file mode 100644
index 000000000000..bdc3b5b06316
--- /dev/null
+++ b/sdk/ledger-test-tool/suites/src/main/scala/com/daml/ledger/api/testtool/suites/v1_17/PrefetchContractKeysIT.scala
@@ -0,0 +1,82 @@
+// Copyright (c) 2024 Digital Asset (Switzerland) GmbH and/or its affiliates. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package com.daml.ledger.api.testtool.suites.v1_17
+
+import com.daml.ledger.api.testtool.infrastructure.Allocation.{
+  Participant,
+  Participants,
+  SingleParty,
+  allocate,
+}
+import com.daml.ledger.api.testtool.infrastructure.LedgerTestSuite
+import com.daml.ledger.api.testtool.suites.v1_8.CompanionImplicits
+import com.daml.ledger.api.v1.commands.PrefetchContractKey
+import com.daml.ledger.javaapi.data.codegen.ByKey
+import com.daml.ledger.test.java.model.da.types.Tuple2
+import com.daml.ledger.test.java.model.test.{Dummy, TextKey, TextKeyOperations, WithKey}
+
+import scala.jdk.CollectionConverters._
+import scala.jdk.OptionConverters._
+
+object PrefetchContractKeysIT {
+  implicit final class JavaBindingSupportExtension(val byKey: ByKey) extends AnyVal {
+    def toPrefetchKey(): PrefetchContractKey = {
+      ???
+    }
+  }
+}
+
+class PrefetchContractKeysIT extends LedgerTestSuite {
+  import CompanionImplicits._
+  import PrefetchContractKeysIT._
+
+  test(
+    "CSprefetchContractKeysBasic",
+    "Explicit contract key prefetches are accepted",
+    allocate(SingleParty),
+  )(implicit ec => { case Participants(Participant(ledger, party)) =>
+    val prefetch = WithKey.byKey(party).toPrefetchKey()
+    val request = ledger
+      .submitAndWaitRequest(party, new WithKey(party).create.commands)
+      .update(_.commands.prefetchContractKeys := Seq(prefetch))
+    for {
+      _ <- ledger.submitAndWait(request)
+      active <- ledger.activeContracts(party)
+    } yield {
+      assert(active.size == 1)
+      val dummyTemplateId = active.flatMap(_.templateId.toList).head
+      assert(dummyTemplateId == Dummy.TEMPLATE_ID_WITH_PACKAGE_ID.toV1)
+    }
+  })
+
+  test(
+    "CSprefetchContractKeysMany",
+    "Prefetch many contract keys",
+    allocate(SingleParty),
+  )(implicit ec => { case Participants(Participant(ledger, party)) =>
+    val numPrefetches = 1000
+    val prefetches =
+      (1 to numPrefetches).map(i => TextKey.byKey(new Tuple2(party, s"key$i")).toPrefetchKey())
+    val existingKeyIndex = 10
+    for {
+      textKeyContract <- ledger.create(party, new TextKey(party, "key1", Seq.empty.asJava))
+      textKeyOps <- ledger.create(party, new TextKeyOperations(party))
+      exerciseCommands = (1 to numPrefetches)
+        .flatMap(i =>
+          textKeyOps
+            .exerciseTKOLookup(
+              new Tuple2(party, s"key$i"),
+              Option.when(i == existingKeyIndex)(textKeyContract).toJava,
+            )
+            .commands
+            .asScala
+        )
+        .asJava
+      request = ledger
+        .submitAndWaitRequest(party, exerciseCommands)
+        .update(_.commands.prefetchContractKeys := prefetches)
+      _ <- ledger.submitAndWait(request)
+    } yield ()
+  })
+}