diff --git a/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java index 2fe8788c6df68..c15b0063386f2 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java @@ -4,6 +4,7 @@ import static com.linkedin.metadata.Constants.STATUS_ASPECT_NAME; import static com.linkedin.metadata.entity.ebean.EbeanAspectDao.TX_ISOLATION; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; @@ -12,7 +13,6 @@ import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.reset; -import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; @@ -48,6 +48,7 @@ import com.linkedin.metadata.utils.PegasusUtils; import com.linkedin.mxe.MetadataChangeProposal; import com.linkedin.mxe.SystemMetadata; +import com.linkedin.util.Pair; import io.datahubproject.metadata.context.OperationContext; import io.datahubproject.metadata.context.RetrieverContext; import io.datahubproject.test.DataGenerator; @@ -58,7 +59,6 @@ import jakarta.persistence.EntityNotFoundException; import java.net.URISyntaxException; import java.sql.Timestamp; -import java.util.ArrayList; import java.util.Collection; import java.util.List; import java.util.Map; @@ -90,9 +90,6 @@ public class EbeanEntityServiceTest extends EntityServiceTest { - // Track additional Database instances created in individual tests for cleanup - private final List additionalDatabases = new ArrayList<>(); - public EbeanEntityServiceTest() throws EntityRegistryException {} @DataProvider(name = "cdcModeVariants") @@ -170,11 +167,8 @@ public void testNoRowsUpdatedErrorHandling() throws Exception { CorpUserInfo writeAspect = AspectGenerationUtils.createCorpUserInfo("email@test.com"); String aspectName = PegasusUtils.getAspectNameFromSchema(writeAspect.schema()); - // Create database and spy on aspectDao - Database server = EbeanTestUtils.createTestServer(EbeanEntityServiceTest.class.getSimpleName()); - additionalDatabases.add(server); // Track for cleanup - EbeanAspectDao aspectDao = - spy(new EbeanAspectDao(server, EbeanConfiguration.testDefault, null)); + // No database needed since all methods are stubbed + EbeanAspectDao aspectDao = mock(EbeanAspectDao.class); // Prevent actual saves EntityAspect mockEntityAspect = mock(EntityAspect.class); @@ -182,7 +176,16 @@ public void testNoRowsUpdatedErrorHandling() throws Exception { doReturn(Optional.of(mockEntityAspect)).when(aspectDao).updateAspect(any(), any()); doReturn(Optional.of(mockEntityAspect)).when(aspectDao).insertAspect(any(), any(), anyLong()); - // Create spied transaction context that throws on commitAndContinue + // Stub methods that the transaction block will call + // Use mutable maps because the code calls computeIfAbsent() on them + when(aspectDao.getLatestAspects(any(), any(), anyBoolean())) + .thenReturn(new java.util.HashMap<>()); + when(aspectDao.getNextVersions(any())).thenReturn(new java.util.HashMap<>()); + // Stub saveLatestAspect to return a Pair with the mocked entity aspects + when(aspectDao.saveLatestAspect(any(), any(), any(), any())) + .thenReturn(Pair.of(Optional.of(mockEntityAspect), Optional.of(mockEntityAspect))); + + // Create mocked transaction context that throws on commitAndContinue AtomicReference capturedTxContext = new AtomicReference<>(); AtomicReference> capturedResult = new AtomicReference<>(); @@ -191,7 +194,14 @@ public void testNoRowsUpdatedErrorHandling() throws Exception { Function> block = invocation.getArgument(0); Integer maxTransactionRetry = invocation.getArgument(2); - TransactionContext txContext = spy(TransactionContext.empty(maxTransactionRetry)); + // Use mock instead of spy to avoid Mockito global interceptor + TransactionContext txContext = mock(TransactionContext.class); + // Stub other methods that might be called + when(txContext.getFailedAttempts()).thenReturn(0); + when(txContext.shouldAttemptRetry()).thenReturn(true); + when(txContext.lastException()).thenReturn(null); + when(txContext.lastExceptionIsDuplicateKey()).thenReturn(false); + capturedTxContext.set(txContext); doThrow(new EntityNotFoundException("No rows updated")) @@ -205,7 +215,7 @@ public void testNoRowsUpdatedErrorHandling() throws Exception { .when(aspectDao) .runInTransactionWithRetry(any(), any(), anyInt()); - // Create the service with our spied dao + // Create the service with our mocked dao PreProcessHooks preProcessHooks = new PreProcessHooks(); preProcessHooks.setUiEnabled(false); EntityServiceImpl entityService = @@ -778,7 +788,5 @@ public void cleanup() { // Shutdown all Database instances to prevent thread pool and connection leaks // This includes the "gma.heartBeat" thread and connection pools EbeanTestUtils.shutdownDatabaseFromAspectDao(_aspectDao); - EbeanTestUtils.shutdownDatabases(additionalDatabases); - additionalDatabases.clear(); } }