From bb0c58577459fbe265163dbe82789c0f0196e149 Mon Sep 17 00:00:00 2001 From: David Leifker Date: Wed, 19 Jul 2023 11:07:19 -0500 Subject: [PATCH] fixing tests --- .../upgrade/nocode/DataMigrationStep.java | 6 ++++- .../ebean/transactions/AspectsBatch.java | 2 +- .../token/StatefulTokenService.java | 1 - .../DataHubTokenAuthenticatorTest.java | 1 + .../token/StatefulTokenServiceTest.java | 1 + .../test/resources/test-entity-registry.yaml | 10 +++++++ .../boot/steps/IngestPoliciesStep.java | 6 ++++- .../IngestDataPlatformInstancesStepTest.java | 1 + .../test/resources/test-entity-registry.yaml | 6 +++++ .../java/entities/EntitiesControllerTest.java | 10 +++++++ .../resources/entity/AspectResource.java | 26 +++++++++++++------ .../resources/entity/AspectResourceTest.java | 9 ++++++- 12 files changed, 66 insertions(+), 13 deletions(-) diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/DataMigrationStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/DataMigrationStep.java index fb9272f68c37a1..a3c92c04864154 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/DataMigrationStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/DataMigrationStep.java @@ -21,7 +21,11 @@ import io.ebean.EbeanServer; import io.ebean.PagedList; import java.net.URISyntaxException; -import java.util.*; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; import java.util.concurrent.TimeUnit; import java.util.function.Function; diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java index 7c5416f99563e6..042564fc7eb875 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/ebean/transactions/AspectsBatch.java @@ -64,7 +64,7 @@ private static AspectsBatchItem toAspectBatchItem(MetadataChangeProposal mcp, En Urn urn = mcp.getEntityUrn(); if (urn == null) { - urn = EntityKeyUtils.getUrnFromProposal(mcp, aspectSpec); + urn = EntityKeyUtils.getUrnFromProposal(mcp, entitySpec.getKeyAspectSpec()); } return AspectsBatchItem.builder() diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java b/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java index 6c73c009db9bf4..40cf3ae1d5b9a6 100644 --- a/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java +++ b/metadata-service/auth-impl/src/main/java/com/datahub/authentication/token/StatefulTokenService.java @@ -20,7 +20,6 @@ import java.util.Base64; import java.util.Date; import java.util.HashMap; -import java.util.List; import java.util.Map; import java.util.Objects; import java.util.concurrent.ExecutionException; diff --git a/metadata-service/auth-impl/src/test/java/com/datahub/authentication/authenticator/DataHubTokenAuthenticatorTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authentication/authenticator/DataHubTokenAuthenticatorTest.java index 44673f693b555a..f5ce938c411c6f 100644 --- a/metadata-service/auth-impl/src/test/java/com/datahub/authentication/authenticator/DataHubTokenAuthenticatorTest.java +++ b/metadata-service/auth-impl/src/test/java/com/datahub/authentication/authenticator/DataHubTokenAuthenticatorTest.java @@ -106,6 +106,7 @@ public void testAuthenticateSuccess() throws Exception { final AspectSpec keyAspectSpec = configEntityRegistry.getEntitySpec(Constants.ACCESS_TOKEN_ENTITY_NAME).getKeyAspectSpec(); Mockito.when(mockService.getKeyAspectSpec(Mockito.eq(Constants.ACCESS_TOKEN_ENTITY_NAME))).thenReturn(keyAspectSpec); Mockito.when(mockService.exists(Mockito.any(Urn.class))).thenReturn(true); + Mockito.when(mockService.getEntityRegistry()).thenReturn(configEntityRegistry); final DataHubTokenAuthenticator authenticator = new DataHubTokenAuthenticator(); authenticator.init(ImmutableMap.of(SIGNING_KEY_CONFIG_NAME, TEST_SIGNING_KEY, SALT_CONFIG_NAME, diff --git a/metadata-service/auth-impl/src/test/java/com/datahub/authentication/token/StatefulTokenServiceTest.java b/metadata-service/auth-impl/src/test/java/com/datahub/authentication/token/StatefulTokenServiceTest.java index 99f242bac6b42d..12c1307f996d67 100644 --- a/metadata-service/auth-impl/src/test/java/com/datahub/authentication/token/StatefulTokenServiceTest.java +++ b/metadata-service/auth-impl/src/test/java/com/datahub/authentication/token/StatefulTokenServiceTest.java @@ -159,6 +159,7 @@ public void generateRevokeToken() throws TokenException { DataHubTokenAuthenticatorTest.class.getClassLoader().getResourceAsStream("test-entity-registry.yaml")); final AspectSpec keyAspectSpec = configEntityRegistry.getEntitySpec(Constants.ACCESS_TOKEN_ENTITY_NAME).getKeyAspectSpec(); + Mockito.when(mockService.getEntityRegistry()).thenReturn(configEntityRegistry); Mockito.when(mockService.getKeyAspectSpec(Mockito.eq(Constants.ACCESS_TOKEN_ENTITY_NAME))).thenReturn(keyAspectSpec); Mockito.when(mockService.exists(Mockito.any(Urn.class))).thenReturn(true); final RollbackRunResult result = new RollbackRunResult(ImmutableList.of(), 0); diff --git a/metadata-service/auth-impl/src/test/resources/test-entity-registry.yaml b/metadata-service/auth-impl/src/test/resources/test-entity-registry.yaml index acdc5ead92fb1f..48c1b9450bd7d9 100644 --- a/metadata-service/auth-impl/src/test/resources/test-entity-registry.yaml +++ b/metadata-service/auth-impl/src/test/resources/test-entity-registry.yaml @@ -4,4 +4,14 @@ entities: keyAspect: dataHubAccessTokenKey aspects: - dataHubAccessTokenInfo + - name: corpuser + keyAspect: corpUserKey + aspects: + - corpUserInfo + - corpUserEditableInfo + - corpUserStatus + - groupMembership + - status + - corpUserCredentials + - corpUserSettings events: diff --git a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java index 1de96218332ec3..67291946f83dd8 100644 --- a/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java +++ b/metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java @@ -27,7 +27,11 @@ import com.linkedin.policy.DataHubPolicyInfo; import java.io.IOException; import java.net.URISyntaxException; -import java.util.*; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; diff --git a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDataPlatformInstancesStepTest.java b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDataPlatformInstancesStepTest.java index 68618a9e87a79b..f1950d3b798d94 100644 --- a/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDataPlatformInstancesStepTest.java +++ b/metadata-service/factories/src/test/java/com/linkedin/metadata/boot/steps/IngestDataPlatformInstancesStepTest.java @@ -142,6 +142,7 @@ private void mockDBWithWorkToDo( when(migrationsDao.checkIfAspectExists(DATA_PLATFORM_INSTANCE_ASPECT_NAME)).thenReturn(false); when(migrationsDao.countEntities()).thenReturn((long) allUrnsInDB.size()); when(migrationsDao.listAllUrns(anyInt(), anyInt())).thenReturn(allUrnsInDB); + when(entityService.getEntityRegistry()).thenReturn(entityRegistry); } private List insertMockEntities(int count, String entity, String urnTemplate, EntityRegistry entityRegistry, EntityService entityService) { diff --git a/metadata-service/factories/src/test/resources/test-entity-registry.yaml b/metadata-service/factories/src/test/resources/test-entity-registry.yaml index 45aa9b9554fb4b..fe32b413751e69 100644 --- a/metadata-service/factories/src/test/resources/test-entity-registry.yaml +++ b/metadata-service/factories/src/test/resources/test-entity-registry.yaml @@ -8,3 +8,9 @@ entities: keyAspect: chartKey aspects: - domains + - dataPlatformInstance + - name: dataPlatform + category: core + keyAspect: dataPlatformKey + aspects: + - dataPlatformInfo \ No newline at end of file diff --git a/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java b/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java index 2dff636de2f27c..4cf2758d86f70e 100644 --- a/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java +++ b/metadata-service/openapi-servlet/src/test/java/entities/EntitiesControllerTest.java @@ -9,6 +9,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.linkedin.metadata.config.PreProcessHooks; import com.linkedin.metadata.entity.AspectDao; +import com.linkedin.metadata.entity.EntityService; import com.linkedin.metadata.event.EventProducer; import com.linkedin.metadata.models.registry.EntityRegistry; import com.linkedin.metadata.service.UpdateIndicesService; @@ -35,8 +36,12 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.function.Function; + +import io.ebean.Transaction; import mock.MockEntityRegistry; import mock.MockEntityService; +import org.mockito.ArgumentMatchers; import org.mockito.Mockito; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -60,6 +65,11 @@ public void setup() throws NoSuchMethodException, InvocationTargetException, InstantiationException, IllegalAccessException { EntityRegistry mockEntityRegistry = new MockEntityRegistry(); AspectDao aspectDao = Mockito.mock(AspectDao.class); + Mockito.when(aspectDao.runInTransactionWithRetry( + ArgumentMatchers.>any(), anyInt())).thenAnswer(i -> + ((Function) i.getArgument(0)).apply(Mockito.mock(Transaction.class)) + ); + EventProducer mockEntityEventProducer = Mockito.mock(EventProducer.class); UpdateIndicesService mockUpdateIndicesService = mock(UpdateIndicesService.class); PreProcessHooks preProcessHooks = new PreProcessHooks(); diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/AspectResource.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/AspectResource.java index a90c95767c85fd..cf798b1461ade2 100644 --- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/AspectResource.java +++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/entity/AspectResource.java @@ -43,6 +43,7 @@ import io.opentelemetry.extension.annotations.WithSpan; import java.net.URISyntaxException; import java.time.Clock; +import java.util.List; import java.util.Set; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -184,7 +185,7 @@ public Task ingestProposal( @ActionParam(PARAM_ASYNC) @Optional(UNSET) String async) throws URISyntaxException { log.info("INGEST PROPOSAL proposal: {}", metadataChangeProposal); - boolean asyncBool; + final boolean asyncBool; if (UNSET.equals(async)) { asyncBool = Boolean.parseBoolean(System.getenv(ASYNC_INGEST_DEFAULT_NAME)); } else { @@ -205,12 +206,20 @@ public Task ingestProposal( return RestliUtil.toTask(() -> { log.debug("Proposal: {}", metadataChangeProposal); try { - Stream proposalStream = Stream.concat(Stream.of(metadataChangeProposal), - AspectUtils.getAdditionalChanges(metadataChangeProposal, _entityService).stream()); + final AspectsBatch batch; + if (asyncBool) { + // if async we'll expand the additional changes later, no need to do this early + batch = AspectsBatch.builder() + .mcps(List.of(metadataChangeProposal), _entityService.getEntityRegistry()) + .build(); + } else { + Stream proposalStream = Stream.concat(Stream.of(metadataChangeProposal), + AspectUtils.getAdditionalChanges(metadataChangeProposal, _entityService).stream()); - AspectsBatch batch = AspectsBatch.builder() - .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) - .build(); + batch = AspectsBatch.builder() + .mcps(proposalStream.collect(Collectors.toList()), _entityService.getEntityRegistry()) + .build(); + } Set> results = _entityService.ingestProposal(batch, auditStamp, asyncBool); @@ -221,10 +230,11 @@ public Task ingestProposal( .get(); // Update runIds + Urn resultUrn = one.getUrn(); if (!one.isQueued()) { - tryIndexRunId(urn, metadataChangeProposal.getSystemMetadata(), _entitySearchService); + tryIndexRunId(resultUrn, metadataChangeProposal.getSystemMetadata(), _entitySearchService); } - return urn.toString(); + return resultUrn.toString(); } catch (ValidationException e) { throw new RestLiServiceException(HttpStatus.S_422_UNPROCESSABLE_ENTITY, e.getMessage()); } diff --git a/metadata-service/restli-servlet-impl/src/test/java/com/linkedin/metadata/resources/entity/AspectResourceTest.java b/metadata-service/restli-servlet-impl/src/test/java/com/linkedin/metadata/resources/entity/AspectResourceTest.java index 0d3642b91758a1..12bd0c1628c13a 100644 --- a/metadata-service/restli-servlet-impl/src/test/java/com/linkedin/metadata/resources/entity/AspectResourceTest.java +++ b/metadata-service/restli-servlet-impl/src/test/java/com/linkedin/metadata/resources/entity/AspectResourceTest.java @@ -22,6 +22,8 @@ import com.linkedin.mxe.MetadataChangeLog; import com.linkedin.mxe.MetadataChangeProposal; import java.net.URISyntaxException; +import java.util.List; + import mock.MockEntityRegistry; import org.testng.annotations.BeforeTest; import org.testng.annotations.Test; @@ -77,7 +79,12 @@ public void testAsyncDefaultAspects() throws URISyntaxException { reset(_producer, _aspectDao); when(_aspectDao.runInTransactionWithRetry(any(), anyInt())) - .thenReturn(new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0)); + .thenReturn(List.of( + new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0), + new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0), + new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0), + new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0), + new EntityService.UpdateAspectResult(urn, null, properties, null, null, null, null, 0))); _aspectResource.ingestProposal(mcp, "false"); verify(_producer, times(5)).produceMetadataChangeLog(eq(urn), any(AspectSpec.class), any(MetadataChangeLog.class)); verifyNoMoreInteractions(_producer);