fix(test): adapt tests to new Runtime model in EDC (#204)

* fix(test): adapt tests to new runtime model

* DEPENDENCIES

* DEPENDENCIES

* trigger ci

* DEPENDENCIES

Author: paullatzelsperger

DEPENDENCIES

@@ -10,6 +10,9 @@ dependencies {
 
     testImplementation(libs.edc.sql.core)
     testImplementation(libs.edc.sql.dataplane.instancestore)
+    testImplementation(libs.edc.junit)
+    testImplementation(testFixtures(libs.edc.sql.lease))
+    testImplementation(testFixtures(libs.edc.sql.core))
     testImplementation(testFixtures(libs.edc.dpf.selector.spi))
     testImplementation(testFixtures(project(":extensions:common:azure:azure-test")))
 }

extensions/data-plane-selector/data-plane-instance-store-cosmos/build.gradle.kts

@@ -25,21 +25,29 @@
 import org.eclipse.edc.json.JacksonTypeManager;
 import org.eclipse.edc.policy.model.PolicyRegistrationTypes;
 import org.eclipse.edc.sql.QueryExecutor;
+import org.eclipse.edc.sql.lease.testfixtures.LeaseUtil;
 import org.eclipse.edc.transaction.datasource.spi.DataSourceRegistry;
 import org.eclipse.edc.transaction.spi.TransactionContext;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.extension.ExtendWith;
 
+import java.sql.SQLException;
+import java.time.Clock;
+import java.time.Duration;
+import javax.sql.DataSource;
+
 import static org.eclipse.edc.azure.testfixtures.CosmosPostgresTestExtension.DEFAULT_DATASOURCE_NAME;
 import static org.eclipse.edc.junit.testfixtures.TestUtils.getResourceFileContentAsString;
 
 @ParallelPostgresCosmosTest
 @ExtendWith(CosmosPostgresTestExtension.class)
 public class CosmosDataPlaneInstanceStoreTest extends DataPlaneInstanceStoreTestBase {
     private static final DataPlaneInstanceStatements STATEMENTS = new PostgresDataPlaneInstanceStatements();
-    SqlDataPlaneInstanceStore store;
+    private final DataPlaneInstanceStatements statements = new PostgresDataPlaneInstanceStatements();
+    private SqlDataPlaneInstanceStore store;
+    private LeaseUtil leaseUtil;
 
     @BeforeAll
     static void createDatabase(CosmosPostgresTestExtension.SqlHelper helper) {
@@ -52,13 +60,22 @@ static void dropDatabase(CosmosPostgresTestExtension.SqlHelper helper) {
     }
 
     @BeforeEach
-    void setUp(DataSourceRegistry reg, TransactionContext transactionContext, QueryExecutor queryExecutor, CosmosPostgresTestExtension.SqlHelper helper) {
+    void setUp(DataSourceRegistry reg, DataSource dataSource, TransactionContext transactionContext, QueryExecutor queryExecutor, CosmosPostgresTestExtension.SqlHelper helper) {
+
+        var clock = Clock.systemUTC();
 
         var typeManager = new JacksonTypeManager();
         typeManager.registerTypes(DataPlaneInstance.class);
         typeManager.registerTypes(PolicyRegistrationTypes.TYPES.toArray(Class<?>[]::new));
+        leaseUtil = new LeaseUtil(transactionContext, () -> {
+            try {
+                return dataSource.getConnection();
+            } catch (SQLException e) {
+                throw new RuntimeException(e);
+            }
+        }, statements, clock);
 
-        store = new SqlDataPlaneInstanceStore(reg, DEFAULT_DATASOURCE_NAME, transactionContext, STATEMENTS, typeManager.getMapper(), queryExecutor);
+        store = new SqlDataPlaneInstanceStore(reg, DEFAULT_DATASOURCE_NAME, transactionContext, STATEMENTS, typeManager.getMapper(), queryExecutor, clock, CONNECTOR_NAME);
         helper.truncateTable(STATEMENTS.getDataPlaneInstanceTable());
     }
 
@@ -68,5 +85,18 @@ protected DataPlaneInstanceStore getStore() {
         return store;
     }
 
+    @Override
+    protected void leaseEntity(String entityId, String owner, Duration duration) {
+        leaseUtil.leaseEntity(entityId, owner, duration);
+    }
 
+    @Override
+    protected boolean isLeasedBy(String entityId, String owner) {
+        return leaseUtil.isLeased(entityId, owner);
+    }
+
+    @Override
+    protected Duration getTestTimeout() {
+        return Duration.ofSeconds(2);
+    }
 }

extensions/data-plane-selector/data-plane-instance-store-cosmos/src/test/java/org/eclipse/edc/connector/dataplane/selector/store/cosmos/CosmosDataPlaneInstanceStoreTest.java

@@ -1,5 +1,20 @@
+CREATE TABLE IF NOT EXISTS edc_lease
+(
+    leased_by      VARCHAR NOT NULL,
+    leased_at      BIGINT,
+    lease_duration INTEGER NOT NULL,
+    lease_id       VARCHAR NOT NULL
+        CONSTRAINT lease_pk
+            PRIMARY KEY
+);
+
+
 CREATE TABLE IF NOT EXISTS edc_data_plane_instance
 (
-    id                   VARCHAR NOT NULL PRIMARY KEY,
-    data                 JSON
+    id       VARCHAR NOT NULL PRIMARY KEY,
+    data     JSON,
+    lease_id VARCHAR
+        CONSTRAINT data_plane_instance_lease_id_fk
+            REFERENCES edc_lease
+            ON DELETE SET NULL
 );

extensions/data-plane-selector/data-plane-instance-store-cosmos/src/test/resources/schema.sql

@@ -29,6 +29,7 @@
 import org.eclipse.edc.connector.dataplane.spi.store.DataPlaneStore;
 import org.eclipse.edc.json.JacksonTypeManager;
 import org.eclipse.edc.junit.extensions.EdcExtension;
+import org.eclipse.edc.junit.extensions.EmbeddedRuntime;
 import org.eclipse.edc.junit.testfixtures.TestUtils;
 import org.eclipse.edc.spi.types.TypeManager;
 import org.eclipse.edc.spi.types.domain.DataAddress;
@@ -98,7 +99,7 @@ static void afterAll() {
 
     @Test
     void transfer_success(
-            EdcExtension edc,
+            EmbeddedRuntime edc,
             AzureResourceManager azure,
             DataPlaneManager dataPlaneManager,
             DataPlaneStore store) {
@@ -184,7 +185,7 @@ private static class Account {
         private final BlobServiceClient client;
         private final String containerName = "test-container-" + UUID.randomUUID();
 
-        Account(AzureResourceManager azure, EdcExtension edc, String setting) {
+        Account(AzureResourceManager azure, EmbeddedRuntime edc, String setting) {
             String accountId = Objects.requireNonNull(edc.getContext().getConfig().getString(setting), setting);
             var account = azure.storageAccounts().getById(accountId);
             name = account.name();

extensions/data-plane/data-plane-azure-data-factory/src/test/java/org/eclipse/edc/connector/dataplane/azure/datafactory/AzureDataFactoryCopyIntegrationTest.java

@@ -128,8 +128,8 @@ void transferBlob_success(String assetName, String[] blobsToTransfer) {
         providerClient.createContractDefinition(assetId, UUID.randomUUID().toString(), policyId, policyId);
 
         // Write Key to vault
-        consumer.getContext().getService(Vault.class).storeSecret(format("%s-key1", CONSUMER_STORAGE_ACCOUNT_NAME), CONSUMER_STORAGE_ACCOUNT_KEY);
-        provider.getContext().getService(Vault.class).storeSecret(format("%s-key1", PROVIDER_STORAGE_ACCOUNT_NAME), PROVIDER_STORAGE_ACCOUNT_KEY);
+        consumer.getService(Vault.class).storeSecret(format("%s-key1", CONSUMER_STORAGE_ACCOUNT_NAME), CONSUMER_STORAGE_ACCOUNT_KEY);
+        provider.getService(Vault.class).storeSecret(format("%s-key1", PROVIDER_STORAGE_ACCOUNT_NAME), PROVIDER_STORAGE_ACCOUNT_KEY);
 
         var transferProcessId = consumerClient.requestAssetAndTransferToBlob(providerClient, assetId, CONSUMER_STORAGE_ACCOUNT_NAME);
         await().pollInterval(POLL_INTERVAL).atMost(TIMEOUT).untilAsserted(() -> {