fixing tests

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;
 

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()

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;

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,

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);

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:

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;

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<Urn> insertMockEntities(int count, String entity, String urnTemplate, EntityRegistry entityRegistry, EntityService entityService) {

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

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.<Function<Transaction, EntityService.UpdateAspectResult>>any(), anyInt())).thenAnswer(i ->
+            ((Function<Transaction, EntityService.UpdateAspectResult>) i.getArgument(0)).apply(Mockito.mock(Transaction.class))
+    );
+
     EventProducer mockEntityEventProducer = Mockito.mock(EventProducer.class);
     UpdateIndicesService mockUpdateIndicesService = mock(UpdateIndicesService.class);
     PreProcessHooks preProcessHooks = new PreProcessHooks();

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<String> 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<String> ingestProposal(
     return RestliUtil.toTask(() -> {
       log.debug("Proposal: {}", metadataChangeProposal);
       try {
-        Stream<MetadataChangeProposal> 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<MetadataChangeProposal> 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<Pair<AspectsBatchItem, EntityService.IngestProposalResult>> results =
                 _entityService.ingestProposal(batch, auditStamp, asyncBool);
@@ -221,10 +230,11 @@ public Task<String> 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());
       }

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);