diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/CustomPropertiesPatchBuilder.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/CustomPropertiesPatchBuilder.java index e4143851afbe51..b78d563147e636 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/CustomPropertiesPatchBuilder.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/CustomPropertiesPatchBuilder.java @@ -17,10 +17,16 @@ public class CustomPropertiesPatchBuilder> operations = new ArrayList<>(); + private final List> operations; public CustomPropertiesPatchBuilder(T parentBuilder) { this.parent = parentBuilder; + if (parentBuilder != null) { + // If a parent builder is provided, we use the same path operations list. + this.operations = parentBuilder.getPathValues(); + } else { + this.operations = new ArrayList<>(); + } } /** @@ -72,9 +78,4 @@ public CustomPropertiesPatchBuilder setProperties(Map propert public T getParent() { return parent; } - - @Override - public List> getSubPaths() { - return operations; - } } diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DataFlowInfoPatchBuilder.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DataFlowInfoPatchBuilder.java index 6a114d90875fe3..231956a2fcec81 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DataFlowInfoPatchBuilder.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DataFlowInfoPatchBuilder.java @@ -4,12 +4,10 @@ import static com.linkedin.metadata.Constants.DATA_FLOW_ENTITY_NAME; import static com.linkedin.metadata.Constants.DATA_FLOW_INFO_ASPECT_NAME; -import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ObjectNode; import com.linkedin.common.TimeStamp; import com.linkedin.metadata.aspect.patch.PatchOperationType; import com.linkedin.metadata.aspect.patch.builder.subtypesupport.CustomPropertiesPatchBuilderSupport; -import java.util.List; import java.util.Map; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -87,28 +85,23 @@ public DataFlowInfoPatchBuilder setCreated(@Nullable TimeStamp created) { } public DataFlowInfoPatchBuilder setLastModified(@Nullable TimeStamp lastModified) { + ObjectNode lastModifiedNode = instance.objectNode(); if (lastModified == null) { pathValues.add( ImmutableTriple.of( PatchOperationType.REMOVE.getValue(), BASE_PATH + LAST_MODIFIED_KEY, null)); + } else { + lastModifiedNode.put(TIME_KEY, lastModified.getTime()); + if (lastModified.getActor() != null) { + lastModifiedNode.put(ACTOR_KEY, lastModified.getActor().toString()); + } + pathValues.add( + ImmutableTriple.of( + PatchOperationType.ADD.getValue(), BASE_PATH + LAST_MODIFIED_KEY, lastModifiedNode)); } - ObjectNode lastModifiedNode = instance.objectNode(); - lastModifiedNode.put(TIME_KEY, lastModified.getTime()); - if (lastModified.getActor() != null) { - lastModifiedNode.put(ACTOR_KEY, lastModified.getActor().toString()); - } - pathValues.add( - ImmutableTriple.of( - PatchOperationType.ADD.getValue(), BASE_PATH + LAST_MODIFIED_KEY, lastModifiedNode)); return this; } - @Override - protected List> getPathValues() { - pathValues.addAll(customPropertiesPatchBuilder.getSubPaths()); - return pathValues; - } - @Override protected String getAspectName() { return DATA_FLOW_INFO_ASPECT_NAME; diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DataJobInfoPatchBuilder.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DataJobInfoPatchBuilder.java index 99c0ac6c15eb1a..dd17fbacf338eb 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DataJobInfoPatchBuilder.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DataJobInfoPatchBuilder.java @@ -4,13 +4,11 @@ import static com.linkedin.metadata.Constants.DATA_JOB_ENTITY_NAME; import static com.linkedin.metadata.Constants.DATA_JOB_INFO_ASPECT_NAME; -import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ObjectNode; import com.linkedin.common.TimeStamp; import com.linkedin.common.urn.DataFlowUrn; import com.linkedin.metadata.aspect.patch.PatchOperationType; import com.linkedin.metadata.aspect.patch.builder.subtypesupport.CustomPropertiesPatchBuilderSupport; -import java.util.List; import java.util.Map; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -113,12 +111,6 @@ public DataJobInfoPatchBuilder setLastModified(@Nullable TimeStamp lastModified) return this; } - @Override - protected List> getPathValues() { - pathValues.addAll(customPropertiesPatchBuilder.getSubPaths()); - return pathValues; - } - @Override protected String getAspectName() { return DATA_JOB_INFO_ASPECT_NAME; diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DatasetPropertiesPatchBuilder.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DatasetPropertiesPatchBuilder.java index 31e181fc244fba..60d52c7c720881 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DatasetPropertiesPatchBuilder.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DatasetPropertiesPatchBuilder.java @@ -4,10 +4,8 @@ import static com.linkedin.metadata.Constants.DATASET_ENTITY_NAME; import static com.linkedin.metadata.Constants.DATASET_PROPERTIES_ASPECT_NAME; -import com.fasterxml.jackson.databind.JsonNode; import com.linkedin.metadata.aspect.patch.PatchOperationType; import com.linkedin.metadata.aspect.patch.builder.subtypesupport.CustomPropertiesPatchBuilderSupport; -import java.util.List; import java.util.Map; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -116,12 +114,6 @@ public DatasetPropertiesPatchBuilder setCustomProperties(Map pro return this; } - @Override - protected List> getPathValues() { - pathValues.addAll(customPropertiesPatchBuilder.getSubPaths()); - return pathValues; - } - @Override protected String getAspectName() { return DATASET_PROPERTIES_ASPECT_NAME; diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/subtypesupport/IntermediatePatchBuilder.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/subtypesupport/IntermediatePatchBuilder.java index d891a6b9673da0..cd74818c24e191 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/subtypesupport/IntermediatePatchBuilder.java +++ b/entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/subtypesupport/IntermediatePatchBuilder.java @@ -1,9 +1,6 @@ package com.linkedin.metadata.aspect.patch.builder.subtypesupport; -import com.fasterxml.jackson.databind.JsonNode; import com.linkedin.metadata.aspect.patch.builder.AbstractMultiFieldPatchBuilder; -import java.util.List; -import org.apache.commons.lang3.tuple.ImmutableTriple; /** * Used for supporting intermediate subtypes when constructing a patch for an aspect that includes @@ -15,10 +12,4 @@ public interface IntermediatePatchBuilder> getSubPaths(); } diff --git a/entity-registry/src/test/java/com/linkedin/metadata/aspect/patch/builder/DataFlowInfoPatchBuilderTest.java b/entity-registry/src/test/java/com/linkedin/metadata/aspect/patch/builder/DataFlowInfoPatchBuilderTest.java new file mode 100644 index 00000000000000..612282b7c0238c --- /dev/null +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/patch/builder/DataFlowInfoPatchBuilderTest.java @@ -0,0 +1,280 @@ +package com.linkedin.metadata.aspect.patch.builder; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.linkedin.common.TimeStamp; +import com.linkedin.common.urn.Urn; +import java.net.URISyntaxException; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.commons.lang3.tuple.ImmutableTriple; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +public class DataFlowInfoPatchBuilderTest { + + private TestableDataFlowInfoPatchBuilder builder; + private static final String TEST_URN = "urn:li:dataFlow:(test,flow1,PROD)"; + + // Test helper class to expose protected method + private static class TestableDataFlowInfoPatchBuilder extends DataFlowInfoPatchBuilder { + public List> getTestPathValues() { + return getPathValues(); + } + } + + @BeforeMethod + public void setup() throws URISyntaxException { + builder = new TestableDataFlowInfoPatchBuilder(); + builder.urn(Urn.createFromString(TEST_URN)); + } + + @Test + public void testBuildDoesNotAffectPathValues() throws URISyntaxException { + String testName = "testFlow"; + String testDescription = "Test description"; + + builder.setName(testName).setDescription(testDescription).addCustomProperty("key1", "value1"); + + // First call build() + builder.build(); + + // Then verify we can still access pathValues and they're correct + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 3); + + // Verify the operations are still intact + assertEquals(pathValues.get(0).getLeft(), "add"); + assertEquals(pathValues.get(0).getMiddle(), "/name"); + assertEquals(pathValues.get(0).getRight().asText(), testName); + + assertEquals(pathValues.get(1).getLeft(), "add"); + assertEquals(pathValues.get(1).getMiddle(), "/description"); + assertEquals(pathValues.get(1).getRight().asText(), testDescription); + + assertEquals(pathValues.get(2).getLeft(), "add"); + assertTrue(pathValues.get(2).getMiddle().startsWith("/customProperties/")); + assertEquals(pathValues.get(2).getRight().asText(), "value1"); + + // Verify we can call build() again without issues + builder.build(); + + // And verify pathValues are still accessible and correct + pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 3); + } + + @Test + public void testSetName() { + String testName = "testFlow"; + builder.setName(testName); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "add"); + assertEquals(operation.getMiddle(), "/name"); + assertEquals(operation.getRight().asText(), testName); + } + + @Test + public void testSetDescription() { + String testDescription = "Test description"; + builder.setDescription(testDescription); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "add"); + assertEquals(operation.getMiddle(), "/description"); + assertEquals(operation.getRight().asText(), testDescription); + } + + @Test + public void testSetDescriptionNull() { + builder.setDescription(null); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "remove"); + assertEquals(operation.getMiddle(), "/description"); + assertNull(operation.getRight()); + } + + @Test + public void testSetProject() { + String testProject = "testProject"; + builder.setProject(testProject); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "add"); + assertEquals(operation.getMiddle(), "/project"); + assertEquals(operation.getRight().asText(), testProject); + } + + @Test + public void testSetProjectNull() { + builder.setProject(null); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "remove"); + assertEquals(operation.getMiddle(), "/project"); + assertNull(operation.getRight()); + } + + @Test + public void testSetCreated() throws URISyntaxException { + long time = System.currentTimeMillis(); + String actor = "urn:li:corpuser:testUser"; + TimeStamp created = new TimeStamp(); + created.setTime(time); + created.setActor(Urn.createFromString(actor)); + + builder.setCreated(created); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "add"); + assertEquals(operation.getMiddle(), "/created"); + JsonNode createdNode = operation.getRight(); + assertTrue(createdNode.isObject()); + assertEquals(createdNode.get("time").asLong(), time); + assertEquals(createdNode.get("actor").asText(), actor); + } + + @Test + public void testSetCreatedNull() { + builder.setCreated(null); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "remove"); + assertEquals(operation.getMiddle(), "/created"); + assertNull(operation.getRight()); + } + + @Test + public void testSetLastModified() throws URISyntaxException { + long time = System.currentTimeMillis(); + String actor = "urn:li:corpuser:testUser"; + TimeStamp lastModified = new TimeStamp(); + lastModified.setTime(time); + lastModified.setActor(Urn.createFromString(actor)); + + builder.setLastModified(lastModified); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "add"); + assertEquals(operation.getMiddle(), "/lastModified"); + JsonNode lastModifiedNode = operation.getRight(); + assertTrue(lastModifiedNode.isObject()); + assertEquals(lastModifiedNode.get("time").asLong(), time); + assertEquals(lastModifiedNode.get("actor").asText(), actor); + } + + @Test + public void testSetLastModifiedNull() { + builder.setLastModified(null); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "remove"); + assertEquals(operation.getMiddle(), "/lastModified"); + assertNull(operation.getRight()); + } + + @Test + public void testAddCustomProperties() { + builder.addCustomProperty("key1", "value1").addCustomProperty("key2", "value2"); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 2); + + pathValues.forEach( + operation -> { + assertEquals(operation.getLeft(), "add"); + assertTrue(operation.getMiddle().startsWith("/customProperties/")); + assertTrue(operation.getRight().isTextual()); + }); + } + + @Test + public void testRemoveCustomProperty() { + builder.removeCustomProperty("key1"); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "remove"); + assertEquals(operation.getMiddle(), "/customProperties/key1"); + assertNull(operation.getRight()); + } + + @Test + public void testSetCustomProperties() { + Map properties = new HashMap<>(); + properties.put("key1", "value1"); + properties.put("key2", "value2"); + + builder.setCustomProperties(properties); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "add"); + assertEquals(operation.getMiddle(), "/customProperties"); + assertTrue(operation.getRight().isObject()); + } +} diff --git a/entity-registry/src/test/java/com/linkedin/metadata/aspect/patch/builder/DataJobInputOutputPatchBuilderTest.java b/entity-registry/src/test/java/com/linkedin/metadata/aspect/patch/builder/DataJobInputOutputPatchBuilderTest.java new file mode 100644 index 00000000000000..dc141863e24438 --- /dev/null +++ b/entity-registry/src/test/java/com/linkedin/metadata/aspect/patch/builder/DataJobInputOutputPatchBuilderTest.java @@ -0,0 +1,237 @@ +package com.linkedin.metadata.aspect.patch.builder; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.linkedin.common.Edge; +import com.linkedin.common.urn.DataJobUrn; +import com.linkedin.common.urn.DatasetUrn; +import com.linkedin.common.urn.Urn; +import com.linkedin.metadata.graph.LineageDirection; +import java.net.URISyntaxException; +import java.util.List; +import org.apache.commons.lang3.tuple.ImmutableTriple; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +public class DataJobInputOutputPatchBuilderTest { + + private TestableDataJobInputOutputPatchBuilder builder; + private static final String TEST_DATAJOB_URN = + "urn:li:dataJob:(urn:li:dataFlow:(test,flow1,PROD),job1)"; + private static final String TEST_DATASET_URN = + "urn:li:dataset:(urn:li:dataPlatform:hive,SampleTable,PROD)"; + private static final String TEST_DATASET_FIELD_URN = + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:hive,SampleTable,PROD),id)"; + + // Test helper class to expose protected method + private static class TestableDataJobInputOutputPatchBuilder + extends DataJobInputOutputPatchBuilder { + public List> getTestPathValues() { + return getPathValues(); + } + } + + @BeforeMethod + public void setup() throws URISyntaxException { + builder = new TestableDataJobInputOutputPatchBuilder(); + builder.urn(Urn.createFromString(TEST_DATAJOB_URN)); + } + + @Test + public void testBuildDoesNotAffectPathValues() throws URISyntaxException { + DataJobUrn dataJobUrn = DataJobUrn.createFromString(TEST_DATAJOB_URN); + DatasetUrn datasetUrn = DatasetUrn.createFromString(TEST_DATASET_URN); + + builder + .addInputDatajobEdge(dataJobUrn) + .addInputDatasetEdge(datasetUrn) + .addOutputDatasetEdge(datasetUrn); + + // First call build() + builder.build(); + + // Then verify we can still access pathValues and they're correct + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 3); + + // Verify we can call build() again without issues + builder.build(); + + // And verify pathValues are still accessible and correct + pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 3); + } + + @Test + public void testAddInputDatajobEdge() throws URISyntaxException { + DataJobUrn dataJobUrn = DataJobUrn.createFromString(TEST_DATAJOB_URN); + builder.addInputDatajobEdge(dataJobUrn); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "add"); + assertTrue(operation.getMiddle().startsWith("/inputDatajobEdges/")); + assertTrue(operation.getRight().isObject()); + assertEquals(operation.getRight().get("destinationUrn").asText(), dataJobUrn.toString()); + } + + @Test + public void testRemoveInputDatajobEdge() throws URISyntaxException { + DataJobUrn dataJobUrn = DataJobUrn.createFromString(TEST_DATAJOB_URN); + builder.removeInputDatajobEdge(dataJobUrn); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "remove"); + assertTrue(operation.getMiddle().startsWith("/inputDatajobEdges/")); + assertNull(operation.getRight()); + } + + @Test + public void testAddInputDatasetEdge() throws URISyntaxException { + DatasetUrn datasetUrn = DatasetUrn.createFromString(TEST_DATASET_URN); + builder.addInputDatasetEdge(datasetUrn); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "add"); + assertTrue(operation.getMiddle().startsWith("/inputDatasetEdges/")); + assertTrue(operation.getRight().isObject()); + assertEquals(operation.getRight().get("destinationUrn").asText(), datasetUrn.toString()); + } + + @Test + public void testRemoveInputDatasetEdge() throws URISyntaxException { + DatasetUrn datasetUrn = DatasetUrn.createFromString(TEST_DATASET_URN); + builder.removeInputDatasetEdge(datasetUrn); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "remove"); + assertTrue(operation.getMiddle().startsWith("/inputDatasetEdges/")); + assertNull(operation.getRight()); + } + + @Test + public void testAddOutputDatasetEdge() throws URISyntaxException { + DatasetUrn datasetUrn = DatasetUrn.createFromString(TEST_DATASET_URN); + builder.addOutputDatasetEdge(datasetUrn); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "add"); + assertTrue(operation.getMiddle().startsWith("/outputDatasetEdges/")); + assertTrue(operation.getRight().isObject()); + assertEquals(operation.getRight().get("destinationUrn").asText(), datasetUrn.toString()); + } + + @Test + public void testAddInputDatasetField() throws URISyntaxException { + Urn fieldUrn = Urn.createFromString(TEST_DATASET_FIELD_URN); + builder.addInputDatasetField(fieldUrn); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "add"); + assertTrue(operation.getMiddle().startsWith("/inputDatasetFields/")); + assertTrue(operation.getRight().isTextual()); + assertEquals(operation.getRight().asText(), fieldUrn.toString()); + } + + @Test + public void testRemoveInputDatasetField() throws URISyntaxException { + Urn fieldUrn = Urn.createFromString(TEST_DATASET_FIELD_URN); + builder.removeInputDatasetField(fieldUrn); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "remove"); + assertTrue(operation.getMiddle().startsWith("/inputDatasetFields/")); + assertNull(operation.getRight()); + } + + @Test + public void testAddOutputDatasetField() throws URISyntaxException { + Urn fieldUrn = Urn.createFromString(TEST_DATASET_FIELD_URN); + builder.addOutputDatasetField(fieldUrn); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "add"); + assertTrue(operation.getMiddle().startsWith("/outputDatasetFields/")); + assertTrue(operation.getRight().isTextual()); + assertEquals(operation.getRight().asText(), fieldUrn.toString()); + } + + @Test + public void testAddEdgeWithDirection() throws URISyntaxException { + DatasetUrn datasetUrn = DatasetUrn.createFromString(TEST_DATASET_URN); + Edge edge = new Edge(); + edge.setDestinationUrn(datasetUrn); + + builder.addEdge(edge, LineageDirection.UPSTREAM); + builder.build(); + + List> pathValues = builder.getTestPathValues(); + assertNotNull(pathValues); + assertEquals(pathValues.size(), 1); + + ImmutableTriple operation = pathValues.get(0); + assertEquals(operation.getLeft(), "add"); + assertTrue(operation.getMiddle().startsWith("/inputDatasetEdges/")); + assertTrue(operation.getRight().isObject()); + assertEquals(operation.getRight().get("destinationUrn").asText(), datasetUrn.toString()); + } + + @Test + public void testInvalidEntityTypeThrowsException() throws URISyntaxException { + Urn invalidUrn = Urn.createFromString("urn:li:glossaryTerm:invalid"); + Edge edge = new Edge(); + edge.setDestinationUrn(invalidUrn); + + assertThrows( + IllegalArgumentException.class, + () -> { + builder.addEdge(edge, LineageDirection.UPSTREAM); + }); + } +}