Skip to content

Commit 1c47ddc

Browse files
gdatahubgdatahub
committed
test(aspect migration mutator): Automated tests for implementations of Abstract AspectMigrationMutators
1 parent fecfe25 commit 1c47ddc

3 files changed

Lines changed: 470 additions & 0 deletions

File tree

‎entity-registry/src/test/java/com/linkedin/metadata/aspect/hooks/AspectMigrationMutatorBaseTest.java‎

Lines changed: 234 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,234 @@
1+
package com.linkedin.metadata.aspect.hooks;
2+
3+
import static org.mockito.Mockito.mock;
4+
import static org.mockito.Mockito.when;
5+
import static org.testng.Assert.assertEquals;
6+
import static org.testng.Assert.assertFalse;
7+
import static org.testng.Assert.assertTrue;
8+
9+
import com.linkedin.common.urn.Urn;
10+
import com.linkedin.common.urn.UrnUtils;
11+
import com.linkedin.data.template.RecordTemplate;
12+
import com.linkedin.events.metadata.ChangeType;
13+
import com.linkedin.metadata.aspect.AspectRetriever;
14+
import com.linkedin.metadata.aspect.ReadItem;
15+
import com.linkedin.metadata.aspect.RetrieverContext;
16+
import com.linkedin.metadata.aspect.batch.ChangeMCP;
17+
import com.linkedin.metadata.models.registry.EntityRegistry;
18+
import com.linkedin.mxe.SystemMetadata;
19+
import com.linkedin.test.metadata.aspect.TestEntityRegistry;
20+
import com.linkedin.test.metadata.aspect.batch.TestMCP;
21+
import com.linkedin.util.Pair;
22+
import java.util.List;
23+
import java.util.stream.Collectors;
24+
import javax.annotation.Nonnull;
25+
import org.testng.annotations.BeforeMethod;
26+
import org.testng.annotations.Test;
27+
28+
/**
29+
* Base test class for {@link AspectMigrationMutator} implementations. Enforces the version contract
30+
* and tests the happy path for both write and read mutations, so concrete mutator tests can focus
31+
* on supplying realistic aspect payloads and asserting the correctness of the transformed result.
32+
*
33+
* <p>Extend this class for every concrete {@link AspectMigrationMutator} implementation. All
34+
* contract-level {@code @Test} methods are inherited automatically — the subclass only needs to
35+
* supply three things:
36+
*
37+
* <ol>
38+
* <li>{@link #mutator()} — the implementation under test
39+
* <li>{@link #provideSourceAspect()} — a realistic aspect payload at {@code sourceVersion}
40+
* <li>{@link #assertTransformed(RecordTemplate)} — assertions specific to the migration logic
41+
* </ol>
42+
*
43+
* <p>Example subclass:
44+
*
45+
* <pre>{@code
46+
* public class MyAspectV1ToV2MigratorTest extends AspectMigrationMutatorBaseTest {
47+
*
48+
* @Override protected AspectMigrationMutator mutator() {
49+
* return new MyAspectV1ToV2Migrator();
50+
* }
51+
*
52+
* @Override protected RecordTemplate provideSourceAspect() {
53+
* Owner v1 = new Owner();
54+
* v1.setOldField("someValue");
55+
* return v1;
56+
* }
57+
*
58+
* @Override protected void assertTransformed(RecordTemplate result) {
59+
* Owner migrated = (Owner) result;
60+
* assertEquals(migrated.getNewField(), "expected value");
61+
* assertFalse(migrated.hasOldField());
62+
* }
63+
* }
64+
* }</pre>
65+
*/
66+
public abstract class AspectMigrationMutatorBaseTest {
67+
68+
private static final Urn DATASET_URN =
69+
UrnUtils.getUrn("urn:li:dataset:(urn:li:dataPlatform:hive,test,PROD)");
70+
71+
// Abstract
72+
73+
/** Returns the {@link AspectMigrationMutator} implementation under test. */
74+
@Nonnull
75+
protected abstract AspectMigrationMutator mutator();
76+
77+
/**
78+
* Returns a realistic aspect payload that is at {@link AspectMigrationMutator#getSourceVersion()}
79+
* and ready to be transformed by the mutator.
80+
*/
81+
@Nonnull
82+
protected abstract RecordTemplate provideSourceAspect();
83+
84+
/**
85+
* Asserts that the transformed aspect payload is correct. Called after a successful write or read
86+
* mutation at {@code sourceVersion}.
87+
*
88+
* @param result the aspect payload after {@code transform()} has been applied (read from the item
89+
* via {@link ChangeMCP#getRecordTemplate()})
90+
*/
91+
protected abstract void assertTransformed(@Nonnull RecordTemplate result);
92+
93+
// Shared
94+
95+
protected EntityRegistry entityRegistry;
96+
protected RetrieverContext retrieverContext;
97+
98+
@BeforeMethod
99+
public void setUp() {
100+
entityRegistry = new TestEntityRegistry();
101+
AspectRetriever mockRetriever = mock(AspectRetriever.class);
102+
when(mockRetriever.getEntityRegistry()).thenReturn(entityRegistry);
103+
retrieverContext = mock(RetrieverContext.class);
104+
when(retrieverContext.getAspectRetriever()).thenReturn(mockRetriever);
105+
}
106+
107+
// Version contract (auto-inherited)
108+
109+
@Test
110+
public void contract_targetVersionIsSourceVersionPlusOne() {
111+
assertEquals(
112+
mutator().getTargetVersion(),
113+
mutator().getSourceVersion() + 1,
114+
"targetVersion must equal sourceVersion + 1");
115+
}
116+
117+
@Test
118+
public void contract_sourceVersionAtLeastDefault() {
119+
assertTrue(
120+
mutator().getSourceVersion() >= AspectMigrationMutator.DEFAULT_SCHEMA_VERSION,
121+
"sourceVersion must be >= DEFAULT_SCHEMA_VERSION");
122+
}
123+
124+
// Write path (auto-inherited)
125+
126+
@Test
127+
public void writeMutation_atSourceVersion_mutatesAndBumpsVersion() {
128+
SystemMetadata sm = new SystemMetadata();
129+
sm.setSchemaVersion(mutator().getSourceVersion());
130+
131+
ChangeMCP item = buildItem(provideSourceAspect(), sm);
132+
133+
List<Pair<ChangeMCP, Boolean>> results =
134+
mutator().writeMutation(List.of(item), retrieverContext).collect(Collectors.toList());
135+
136+
assertTrue(results.get(0).getSecond(), "Expected mutated=true at sourceVersion");
137+
assertEquals(
138+
(long) item.getSystemMetadata().getSchemaVersion(),
139+
mutator().getTargetVersion(),
140+
"schemaVersion must be bumped to targetVersion");
141+
142+
assertTransformed(item.getRecordTemplate());
143+
}
144+
145+
@Test
146+
public void writeMutation_atDefaultSchemaVersion_mutatesAndBumpsVersion() {
147+
// null schemaVersion is treated as DEFAULT_SCHEMA_VERSION — only relevant when sourceVersion==1
148+
if (mutator().getSourceVersion() != AspectMigrationMutator.DEFAULT_SCHEMA_VERSION) {
149+
return;
150+
}
151+
ChangeMCP item = buildItem(provideSourceAspect(), new SystemMetadata()); // no schemaVersion set
152+
153+
List<Pair<ChangeMCP, Boolean>> results =
154+
mutator().writeMutation(List.of(item), retrieverContext).collect(Collectors.toList());
155+
156+
assertTrue(results.get(0).getSecond(), "Expected mutated=true for null schemaVersion (=1)");
157+
assertEquals((long) item.getSystemMetadata().getSchemaVersion(), mutator().getTargetVersion());
158+
}
159+
160+
@Test
161+
public void writeMutation_alreadyAtTargetVersion_isNoOp() {
162+
SystemMetadata sm = new SystemMetadata();
163+
sm.setSchemaVersion(mutator().getTargetVersion());
164+
165+
ChangeMCP item = buildItem(provideSourceAspect(), sm);
166+
167+
List<Pair<ChangeMCP, Boolean>> results =
168+
mutator().writeMutation(List.of(item), retrieverContext).collect(Collectors.toList());
169+
170+
assertFalse(results.get(0).getSecond(), "Expected no-op when already at targetVersion");
171+
assertEquals(
172+
(long) item.getSystemMetadata().getSchemaVersion(),
173+
mutator().getTargetVersion(),
174+
"schemaVersion must remain unchanged");
175+
}
176+
177+
@Test
178+
public void writeMutation_futureVersion_isNoOp() {
179+
SystemMetadata sm = new SystemMetadata();
180+
sm.setSchemaVersion(mutator().getTargetVersion() + 1); // ahead of this mutator
181+
182+
ChangeMCP item = buildItem(provideSourceAspect(), sm);
183+
184+
List<Pair<ChangeMCP, Boolean>> results =
185+
mutator().writeMutation(List.of(item), retrieverContext).collect(Collectors.toList());
186+
187+
assertFalse(results.get(0).getSecond(), "Expected no-op for version ahead of targetVersion");
188+
}
189+
190+
// Read path (auto-inherited)
191+
@Test
192+
public void readMutation_atSourceVersion_mutatesAndBumpsVersion() {
193+
SystemMetadata sm = new SystemMetadata();
194+
sm.setSchemaVersion(mutator().getSourceVersion());
195+
196+
ChangeMCP item = buildItem(provideSourceAspect(), sm);
197+
198+
List<Pair<ReadItem, Boolean>> results =
199+
mutator().readMutation(List.of(item), retrieverContext).collect(Collectors.toList());
200+
201+
assertTrue(results.get(0).getSecond(), "Expected mutated=true on read path at sourceVersion");
202+
assertEquals(
203+
(long) item.getSystemMetadata().getSchemaVersion(),
204+
mutator().getTargetVersion(),
205+
"schemaVersion must be bumped on read path");
206+
207+
assertTransformed(item.getRecordTemplate());
208+
}
209+
210+
@Test
211+
public void readMutation_alreadyAtTargetVersion_isNoOp() {
212+
SystemMetadata sm = new SystemMetadata();
213+
sm.setSchemaVersion(mutator().getTargetVersion());
214+
215+
ChangeMCP item = buildItem(provideSourceAspect(), sm);
216+
217+
List<Pair<ReadItem, Boolean>> results =
218+
mutator().readMutation(List.of(item), retrieverContext).collect(Collectors.toList());
219+
220+
assertFalse(results.get(0).getSecond(), "Expected no-op on read path when already migrated");
221+
}
222+
223+
// Helper
224+
protected ChangeMCP buildItem(RecordTemplate record, SystemMetadata sm) {
225+
return TestMCP.builder()
226+
.urn(DATASET_URN)
227+
.changeType(ChangeType.UPSERT)
228+
.entitySpec(entityRegistry.getEntitySpec("dataset"))
229+
.aspectSpec(entityRegistry.getAspectSpecs().get(mutator().getAspectName()))
230+
.recordTemplate(record)
231+
.systemMetadata(sm)
232+
.build();
233+
}
234+
}

‎entity-registry/src/test/java/com/linkedin/metadata/aspect/hooks/AspectMigrationMutatorChainTest.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,10 @@ static class TestOwnershipMutator extends AspectMigrationMutator {
248248
private final long sourceVersion;
249249
private final long targetVersion;
250250

251+
TestOwnershipMutator() {
252+
this(1L, 2L);
253+
}
254+
251255
TestOwnershipMutator(long sourceVersion, long targetVersion) {
252256
this.sourceVersion = sourceVersion;
253257
this.targetVersion = targetVersion;

0 commit comments

Comments
 (0)