44import static com .linkedin .metadata .Constants .STATUS_ASPECT_NAME ;
55import static com .linkedin .metadata .entity .ebean .EbeanAspectDao .TX_ISOLATION ;
66import static org .mockito .ArgumentMatchers .any ;
7+ import static org .mockito .ArgumentMatchers .anyBoolean ;
78import static org .mockito .ArgumentMatchers .anyInt ;
89import static org .mockito .ArgumentMatchers .anyLong ;
910import static org .mockito .ArgumentMatchers .eq ;
4849import com .linkedin .metadata .utils .PegasusUtils ;
4950import com .linkedin .mxe .MetadataChangeProposal ;
5051import com .linkedin .mxe .SystemMetadata ;
52+ import com .linkedin .util .Pair ;
5153import io .datahubproject .metadata .context .OperationContext ;
5254import io .datahubproject .metadata .context .RetrieverContext ;
5355import 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