Skip to content

Commit 5b31839

Browse files
fix(test): mock instead of spy
1 parent e85f876 commit 5b31839

File tree

1 file changed

+21
-3
lines changed

1 file changed

+21
-3
lines changed

metadata-io/src/test/java/com/linkedin/metadata/entity/EbeanEntityServiceTest.java

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import static com.linkedin.metadata.Constants.STATUS_ASPECT_NAME;
55
import static com.linkedin.metadata.entity.ebean.EbeanAspectDao.TX_ISOLATION;
66
import static org.mockito.ArgumentMatchers.any;
7+
import static org.mockito.ArgumentMatchers.anyBoolean;
78
import static org.mockito.ArgumentMatchers.anyInt;
89
import static org.mockito.ArgumentMatchers.anyLong;
910
import static org.mockito.ArgumentMatchers.eq;
@@ -48,6 +49,7 @@
4849
import com.linkedin.metadata.utils.PegasusUtils;
4950
import com.linkedin.mxe.MetadataChangeProposal;
5051
import com.linkedin.mxe.SystemMetadata;
52+
import com.linkedin.util.Pair;
5153
import io.datahubproject.metadata.context.OperationContext;
5254
import io.datahubproject.metadata.context.RetrieverContext;
5355
import io.datahubproject.test.DataGenerator;
@@ -182,7 +184,16 @@ public void testNoRowsUpdatedErrorHandling() throws Exception {
182184
doReturn(Optional.of(mockEntityAspect)).when(aspectDao).updateAspect(any(), any());
183185
doReturn(Optional.of(mockEntityAspect)).when(aspectDao).insertAspect(any(), any(), anyLong());
184186

185-
// Create spied transaction context that throws on commitAndContinue
187+
// Stub methods that the transaction block will call
188+
// Use mutable maps because the code calls computeIfAbsent() on them
189+
when(aspectDao.getLatestAspects(any(), any(), anyBoolean()))
190+
.thenReturn(new java.util.HashMap<>());
191+
when(aspectDao.getNextVersions(any())).thenReturn(new java.util.HashMap<>());
192+
// Stub saveLatestAspect to return a Pair with the mocked entity aspects
193+
when(aspectDao.saveLatestAspect(any(), any(), any(), any()))
194+
.thenReturn(Pair.of(Optional.of(mockEntityAspect), Optional.of(mockEntityAspect)));
195+
196+
// Create mocked transaction context that throws on commitAndContinue
186197
AtomicReference<TransactionContext> capturedTxContext = new AtomicReference<>();
187198
AtomicReference<TransactionResult<?>> capturedResult = new AtomicReference<>();
188199

@@ -191,7 +202,14 @@ public void testNoRowsUpdatedErrorHandling() throws Exception {
191202
Function<TransactionContext, TransactionResult<?>> block = invocation.getArgument(0);
192203
Integer maxTransactionRetry = invocation.getArgument(2);
193204

194-
TransactionContext txContext = spy(TransactionContext.empty(maxTransactionRetry));
205+
// Use mock instead of spy to avoid Mockito global interceptor
206+
TransactionContext txContext = mock(TransactionContext.class);
207+
// Stub other methods that might be called
208+
when(txContext.getFailedAttempts()).thenReturn(0);
209+
when(txContext.shouldAttemptRetry()).thenReturn(true);
210+
when(txContext.lastException()).thenReturn(null);
211+
when(txContext.lastExceptionIsDuplicateKey()).thenReturn(false);
212+
195213
capturedTxContext.set(txContext);
196214

197215
doThrow(new EntityNotFoundException("No rows updated"))
@@ -205,7 +223,7 @@ public void testNoRowsUpdatedErrorHandling() throws Exception {
205223
.when(aspectDao)
206224
.runInTransactionWithRetry(any(), any(), anyInt());
207225

208-
// Create the service with our spied dao
226+
// Create the service with our mocked dao
209227
PreProcessHooks preProcessHooks = new PreProcessHooks();
210228
preProcessHooks.setUiEnabled(false);
211229
EntityServiceImpl entityService =

0 commit comments

Comments
 (0)