Skip to content

Commit e2823b4

Browse files
david-leifkerdavid-leifker
committed
more tests
1 parent 08038f9 commit e2823b4

File tree

3 files changed

+304
-17
lines changed

3 files changed

+304
-17
lines changed

‎metadata-io/src/main/java/com/linkedin/metadata/entity/EntityServiceImpl.java

+22-13
Original file line numberDiff line numberDiff line change
@@ -1672,23 +1672,32 @@ public List<RestoreIndicesResult> restoreIndices(
16721672
batch -> {
16731673
long timeSqlQueryMs = System.currentTimeMillis() - startTime;
16741674

1675-
List<SystemAspect> systemAspects =
1676-
EntityUtils.toSystemAspectFromEbeanAspects(
1677-
opContext.getRetrieverContext(), batch.collect(Collectors.toList()), false);
1675+
try {
1676+
List<SystemAspect> systemAspects =
1677+
EntityUtils.toSystemAspectFromEbeanAspects(
1678+
opContext.getRetrieverContext(),
1679+
batch.collect(Collectors.toList()),
1680+
false);
16781681

1679-
RestoreIndicesResult result =
1680-
restoreIndices(opContext, systemAspects, logger, args.createDefaultAspects());
1681-
result.timeSqlQueryMs = timeSqlQueryMs;
1682+
RestoreIndicesResult result =
1683+
restoreIndices(opContext, systemAspects, logger, args.createDefaultAspects());
1684+
result.timeSqlQueryMs = timeSqlQueryMs;
16821685

1683-
logger.accept("Batch completed.");
1684-
try {
1685-
TimeUnit.MILLISECONDS.sleep(args.batchDelayMs);
1686-
} catch (InterruptedException e) {
1687-
throw new RuntimeException(
1688-
"Thread interrupted while sleeping after successful batch migration.");
1686+
logger.accept("Batch completed.");
1687+
try {
1688+
TimeUnit.MILLISECONDS.sleep(args.batchDelayMs);
1689+
} catch (InterruptedException e) {
1690+
throw new RuntimeException(
1691+
"Thread interrupted while sleeping after successful batch migration.");
1692+
}
1693+
1694+
return result;
1695+
} catch (Exception e) {
1696+
log.error("Error processing aspect for restore indices.", e);
1697+
return null;
16891698
}
1690-
return result;
16911699
})
1700+
.filter(Objects::nonNull)
16921701
.collect(Collectors.toList());
16931702
}
16941703
}

‎metadata-io/src/test/java/com/linkedin/metadata/entity/EntityServiceImplTest.java

+280-2
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,12 @@
55
import static com.linkedin.metadata.Constants.UPSTREAM_LINEAGE_ASPECT_NAME;
66
import static com.linkedin.metadata.entity.EntityServiceTest.TEST_AUDIT_STAMP;
77
import static org.mockito.ArgumentMatchers.any;
8+
import static org.mockito.ArgumentMatchers.anyInt;
9+
import static org.mockito.ArgumentMatchers.argThat;
810
import static org.mockito.Mockito.doReturn;
911
import static org.mockito.Mockito.mock;
1012
import static org.mockito.Mockito.never;
13+
import static org.mockito.Mockito.spy;
1114
import static org.mockito.Mockito.times;
1215
import static org.mockito.Mockito.verify;
1316
import static org.mockito.Mockito.when;
@@ -36,9 +39,12 @@
3639
import com.linkedin.metadata.config.PreProcessHooks;
3740
import com.linkedin.metadata.entity.ebean.EbeanAspectV2;
3841
import com.linkedin.metadata.entity.ebean.EbeanSystemAspect;
42+
import com.linkedin.metadata.entity.ebean.PartitionedStream;
3943
import com.linkedin.metadata.entity.ebean.batch.AspectsBatchImpl;
4044
import com.linkedin.metadata.entity.ebean.batch.ChangeItemImpl;
4145
import com.linkedin.metadata.entity.ebean.batch.DeleteItemImpl;
46+
import com.linkedin.metadata.entity.restoreindices.RestoreIndicesArgs;
47+
import com.linkedin.metadata.entity.restoreindices.RestoreIndicesResult;
4248
import com.linkedin.metadata.event.EventProducer;
4349
import com.linkedin.metadata.models.registry.EntityRegistry;
4450
import com.linkedin.metadata.utils.GenericRecordUtils;
@@ -49,11 +55,13 @@
4955
import io.datahubproject.metadata.context.OperationContext;
5056
import io.datahubproject.test.metadata.context.TestOperationContexts;
5157
import java.sql.Timestamp;
58+
import java.util.ArrayList;
5259
import java.util.List;
5360
import java.util.Optional;
5461
import java.util.concurrent.CompletableFuture;
5562
import java.util.concurrent.Future;
5663
import java.util.stream.Stream;
64+
import org.mockito.ArgumentCaptor;
5765
import org.testng.annotations.BeforeMethod;
5866
import org.testng.annotations.Test;
5967

@@ -520,7 +528,7 @@ public void testAspectWithLineageRelationship() {
520528
@Test
521529
public void testIngestTimeseriesProposal() {
522530
// Create a spy of the EntityServiceImpl to track method calls
523-
EntityServiceImpl entityServiceSpy = org.mockito.Mockito.spy(entityService);
531+
EntityServiceImpl entityServiceSpy = spy(entityService);
524532

525533
Urn timeseriesUrn =
526534
UrnUtils.getUrn("urn:li:dataset:(urn:li:dataPlatform:test,timeseriesTest,PROD)");
@@ -586,7 +594,7 @@ public void testIngestTimeseriesProposal() {
586594
@Test
587595
public void testIngestTimeseriesProposalUnsupported() {
588596
// Create a spy of the EntityServiceImpl to track method calls
589-
EntityServiceImpl entityServiceSpy = org.mockito.Mockito.spy(entityService);
597+
EntityServiceImpl entityServiceSpy = spy(entityService);
590598

591599
Urn timeseriesUrn =
592600
UrnUtils.getUrn("urn:li:dataset:(urn:li:dataPlatform:test,timeseriesUnsupportedTest,PROD)");
@@ -616,4 +624,274 @@ public void testIngestTimeseriesProposalUnsupported() {
616624
// Expected
617625
}
618626
}
627+
628+
/**
629+
* Tests an end-to-end scenario with createDefaultAspects=true, ensuring that both the aspects are
630+
* restored to the index and default aspects are created.
631+
*/
632+
@Test
633+
public void testRestoreIndicesEndToEndWithDefaultAspects() throws Exception {
634+
// Setup mock AspectDao
635+
AspectDao mockAspectDao = mock(AspectDao.class);
636+
PartitionedStream<EbeanAspectV2> mockStream = mock(PartitionedStream.class);
637+
638+
// Create test aspects
639+
List<EbeanAspectV2> batch = new ArrayList<>();
640+
641+
// Dataset aspect
642+
EbeanAspectV2 datasetAspect =
643+
new EbeanAspectV2(
644+
"urn:li:dataset:(urn:li:dataPlatform:test,defaultAspectsTest,PROD)",
645+
STATUS_ASPECT_NAME,
646+
0L,
647+
RecordUtils.toJsonString(new Status().setRemoved(false)),
648+
new Timestamp(System.currentTimeMillis()),
649+
TEST_AUDIT_STAMP.getActor().toString(),
650+
null,
651+
RecordUtils.toJsonString(SystemMetadataUtils.createDefaultSystemMetadata()));
652+
batch.add(datasetAspect);
653+
654+
// Setup mock stream
655+
when(mockStream.partition(anyInt())).thenReturn(Stream.of(batch.stream()));
656+
when(mockAspectDao.streamAspectBatches(any())).thenReturn(mockStream);
657+
658+
// Setup mock EventProducer
659+
EventProducer mockEventProducer = mock(EventProducer.class);
660+
when(mockEventProducer.produceMetadataChangeLog(
661+
any(OperationContext.class), any(), any(), any()))
662+
.thenReturn(CompletableFuture.completedFuture(null));
663+
664+
// Create EntityServiceImpl with mocks
665+
EntityServiceImpl entityServiceSpy =
666+
spy(
667+
new EntityServiceImpl(
668+
mockAspectDao, mockEventProducer, false, mock(PreProcessHooks.class), 0, true));
669+
670+
// Mock ingestProposalSync to capture default aspects
671+
ArgumentCaptor<AspectsBatch> batchCaptor = ArgumentCaptor.forClass(AspectsBatch.class);
672+
doReturn(Stream.empty())
673+
.when(entityServiceSpy)
674+
.ingestProposalSync(any(OperationContext.class), batchCaptor.capture());
675+
676+
// Create RestoreIndicesArgs with createDefaultAspects set to true
677+
RestoreIndicesArgs args =
678+
new RestoreIndicesArgs()
679+
.start(0)
680+
.limit(100)
681+
.batchSize(50)
682+
.batchDelayMs(0L)
683+
.createDefaultAspects(true); // Explicitly set to true
684+
685+
// Execute the method under test
686+
List<RestoreIndicesResult> results =
687+
entityServiceSpy.restoreIndices(opContext, args, message -> {});
688+
689+
// Verify results
690+
assertNotNull(results);
691+
assertEquals(1, results.size());
692+
693+
// Verify MCL production
694+
verify(mockEventProducer)
695+
.produceMetadataChangeLog(any(OperationContext.class), any(), any(), any());
696+
697+
// Verify default aspect creation was attempted
698+
verify(entityServiceSpy).ingestProposalSync(any(OperationContext.class), any());
699+
700+
// Verify the captured batch contains aspects
701+
AspectsBatch capturedBatch = batchCaptor.getValue();
702+
assertNotNull(capturedBatch);
703+
assertFalse(capturedBatch.getItems().isEmpty());
704+
705+
// Verify the defaultAspectsCreated count in results
706+
assertTrue(results.get(0).defaultAspectsCreated >= 0);
707+
}
708+
709+
/**
710+
* Tests an end-to-end scenario with createDefaultAspects=false, ensuring that only aspects are
711+
* restored to the index without creating default aspects.
712+
*/
713+
@Test
714+
public void testRestoreIndicesEndToEndWithoutDefaultAspects() throws Exception {
715+
// Setup mock AspectDao
716+
AspectDao mockAspectDao = mock(AspectDao.class);
717+
PartitionedStream<EbeanAspectV2> mockStream = mock(PartitionedStream.class);
718+
719+
// Create test aspects
720+
List<EbeanAspectV2> batch = new ArrayList<>();
721+
722+
// Dataset aspect
723+
EbeanAspectV2 datasetAspect =
724+
new EbeanAspectV2(
725+
"urn:li:dataset:(urn:li:dataPlatform:test,defaultAspectsTest,PROD)",
726+
STATUS_ASPECT_NAME,
727+
0L,
728+
RecordUtils.toJsonString(new Status().setRemoved(false)),
729+
new Timestamp(System.currentTimeMillis()),
730+
TEST_AUDIT_STAMP.getActor().toString(),
731+
null,
732+
RecordUtils.toJsonString(SystemMetadataUtils.createDefaultSystemMetadata()));
733+
batch.add(datasetAspect);
734+
735+
// Setup mock stream
736+
when(mockStream.partition(anyInt())).thenReturn(Stream.of(batch.stream()));
737+
when(mockAspectDao.streamAspectBatches(any())).thenReturn(mockStream);
738+
739+
// Setup mock EventProducer
740+
EventProducer mockEventProducer = mock(EventProducer.class);
741+
when(mockEventProducer.produceMetadataChangeLog(
742+
any(OperationContext.class), any(), any(), any()))
743+
.thenReturn(CompletableFuture.completedFuture(null));
744+
745+
// Create EntityServiceImpl with mocks
746+
EntityServiceImpl entityServiceSpy =
747+
spy(
748+
new EntityServiceImpl(
749+
mockAspectDao, mockEventProducer, false, mock(PreProcessHooks.class), 0, true));
750+
751+
// Simply stub the method without capturing
752+
doReturn(Stream.empty())
753+
.when(entityServiceSpy)
754+
.ingestProposalSync(any(OperationContext.class), any(AspectsBatch.class));
755+
756+
// Create RestoreIndicesArgs with createDefaultAspects set to false
757+
RestoreIndicesArgs args =
758+
new RestoreIndicesArgs()
759+
.start(0)
760+
.limit(100)
761+
.batchSize(50)
762+
.batchDelayMs(0L)
763+
.createDefaultAspects(false); // Explicitly set to false
764+
765+
// Execute the method under test
766+
List<RestoreIndicesResult> results =
767+
entityServiceSpy.restoreIndices(opContext, args, message -> {});
768+
769+
// Verify results
770+
assertNotNull(results);
771+
assertEquals(1, results.size());
772+
773+
// Verify MCL production
774+
verify(mockEventProducer)
775+
.produceMetadataChangeLog(any(OperationContext.class), any(), any(), any());
776+
777+
// Verify default aspect creation was never attempted
778+
verify(entityServiceSpy, never()).ingestProposalSync(any(OperationContext.class), any());
779+
780+
// Don't try to use the capturedBatch since the method should never be called
781+
// Instead, verify the defaultAspectsCreated count in results is 0
782+
assertEquals(0, results.get(0).defaultAspectsCreated);
783+
}
784+
785+
/**
786+
* Tests the continuation behavior of restoreIndices when an exception occurs. It should continue
787+
* processing other aspects even if one fails.
788+
*/
789+
@Test
790+
public void testRestoreIndicesContinuationOnException() throws Exception {
791+
// Setup mock AspectDao
792+
AspectDao mockAspectDao = mock(AspectDao.class);
793+
PartitionedStream<EbeanAspectV2> mockStream = mock(PartitionedStream.class);
794+
795+
// Create test aspects
796+
List<EbeanAspectV2> batch = new ArrayList<>();
797+
798+
// First aspect (will succeed)
799+
EbeanAspectV2 successAspect =
800+
new EbeanAspectV2(
801+
"urn:li:dataset:(urn:li:dataPlatform:test,success,PROD)",
802+
STATUS_ASPECT_NAME,
803+
0L,
804+
RecordUtils.toJsonString(new Status().setRemoved(false)),
805+
new Timestamp(System.currentTimeMillis()),
806+
TEST_AUDIT_STAMP.getActor().toString(),
807+
null,
808+
RecordUtils.toJsonString(SystemMetadataUtils.createDefaultSystemMetadata()));
809+
batch.add(successAspect);
810+
811+
// Second aspect (will fail)
812+
EbeanAspectV2 failAspect =
813+
new EbeanAspectV2(
814+
"urn:li:dataset:(urn:li:dataPlatform:test,fail,PROD)",
815+
STATUS_ASPECT_NAME,
816+
0L,
817+
"INVALID_JSON", // This will cause deserialization to fail
818+
new Timestamp(System.currentTimeMillis()),
819+
TEST_AUDIT_STAMP.getActor().toString(),
820+
null,
821+
RecordUtils.toJsonString(SystemMetadataUtils.createDefaultSystemMetadata()));
822+
batch.add(failAspect);
823+
824+
// Third aspect (will succeed)
825+
EbeanAspectV2 anotherSuccessAspect =
826+
new EbeanAspectV2(
827+
"urn:li:dataset:(urn:li:dataPlatform:test,anotherSuccess,PROD)",
828+
STATUS_ASPECT_NAME,
829+
0L,
830+
RecordUtils.toJsonString(new Status().setRemoved(false)),
831+
new Timestamp(System.currentTimeMillis()),
832+
TEST_AUDIT_STAMP.getActor().toString(),
833+
null,
834+
RecordUtils.toJsonString(SystemMetadataUtils.createDefaultSystemMetadata()));
835+
batch.add(anotherSuccessAspect);
836+
837+
// Setup mock stream to create two separate batches
838+
// This is the key change - we need to separate the failing aspect into its own batch
839+
when(mockStream.partition(anyInt()))
840+
.thenReturn(
841+
Stream.of(
842+
// First batch with success aspect
843+
Stream.of(successAspect),
844+
// Second batch with failing aspect
845+
Stream.of(failAspect),
846+
// Third batch with another success aspect
847+
Stream.of(anotherSuccessAspect)));
848+
849+
when(mockAspectDao.streamAspectBatches(any())).thenReturn(mockStream);
850+
851+
// Setup mock EventProducer
852+
EventProducer mockEventProducer = mock(EventProducer.class);
853+
when(mockEventProducer.produceMetadataChangeLog(
854+
any(OperationContext.class),
855+
argThat(
856+
urn ->
857+
urn.toString()
858+
.contains("success")), // Only succeed for URNs containing "success"
859+
any(),
860+
any()))
861+
.thenReturn(CompletableFuture.completedFuture(null));
862+
863+
// Create EntityServiceImpl with mocks
864+
EntityServiceImpl entityService =
865+
new EntityServiceImpl(
866+
mockAspectDao, mockEventProducer, false, mock(PreProcessHooks.class), 0, true);
867+
868+
// Create RestoreIndicesArgs
869+
RestoreIndicesArgs args =
870+
new RestoreIndicesArgs()
871+
.start(0)
872+
.limit(100)
873+
.batchSize(1) // Process one aspect at a time
874+
.batchDelayMs(0L)
875+
.createDefaultAspects(false); // Avoid additional complexity
876+
877+
// Execute the method under test
878+
List<RestoreIndicesResult> results =
879+
entityService.restoreIndices(opContext, args, message -> {});
880+
881+
// Verify results
882+
assertNotNull(results);
883+
// We should get exactly 2 results - one for each successful batch
884+
assertEquals(2, results.size());
885+
886+
// We expect to see rowsMigrated = 1 in each successful result
887+
assertEquals(1, results.get(0).rowsMigrated);
888+
assertEquals(1, results.get(1).rowsMigrated);
889+
890+
// The failing aspect should not generate a result at all, as the implementation
891+
// filters out null results (batches that throw exceptions)
892+
893+
// Verify total calls to produceMetadataChangeLog (one for each successful aspect)
894+
verify(mockEventProducer, times(2))
895+
.produceMetadataChangeLog(any(OperationContext.class), any(), any(), any());
896+
}
619897
}

‎metadata-service/services/src/main/java/com/linkedin/metadata/entity/restoreindices/RestoreIndicesArgs.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -78,8 +78,8 @@ public RestoreIndicesArgs lePitEpochMs(Long lePitEpochMs) {
7878
return this;
7979
}
8080

81-
public RestoreIndicesArgs createDefaultAspects(@Nullable Boolean readOnly) {
82-
this.createDefaultAspects = readOnly != null ? readOnly : false;
81+
public RestoreIndicesArgs createDefaultAspects(@Nullable Boolean createDefaultAspects) {
82+
this.createDefaultAspects = createDefaultAspects != null ? createDefaultAspects : false;
8383
return this;
8484
}
8585
}

0 commit comments

Comments
 (0)