fix(java-sdk): custom properties patch client (#11984)

entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/CustomPropertiesPatchBuilder.java

@@ -17,10 +17,16 @@ public class CustomPropertiesPatchBuilder<T extends AbstractMultiFieldPatchBuild
   public static final String CUSTOM_PROPERTIES_BASE_PATH = "/customProperties";
 
   private final T parent;
-  private final List<ImmutableTriple<String, String, JsonNode>> operations = new ArrayList<>();
+  private final List<ImmutableTriple<String, String, JsonNode>> 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 class CustomPropertiesPatchBuilder<T extends AbstractMultiFieldPatchBuild
   public T getParent() {
     return parent;
   }
-
-  @Override
-  public List<ImmutableTriple<String, String, JsonNode>> getSubPaths() {
-    return operations;
-  }
 }

entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DataFlowInfoPatchBuilder.java

@@ -4,12 +4,10 @@ import static com.fasterxml.jackson.databind.node.JsonNodeFactory.instance;
 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,12 +85,12 @@ public class DataFlowInfoPatchBuilder
   }
 
   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 {
-    ObjectNode lastModifiedNode = instance.objectNode();
       lastModifiedNode.put(TIME_KEY, lastModified.getTime());
       if (lastModified.getActor() != null) {
         lastModifiedNode.put(ACTOR_KEY, lastModified.getActor().toString());
@@ -100,13 +98,8 @@ public class DataFlowInfoPatchBuilder
       pathValues.add(
           ImmutableTriple.of(
               PatchOperationType.ADD.getValue(), BASE_PATH + LAST_MODIFIED_KEY, lastModifiedNode));
-    return this;
     }
-
+    return this;
-  @Override
-  protected List<ImmutableTriple<String, String, JsonNode>> getPathValues() {
-    pathValues.addAll(customPropertiesPatchBuilder.getSubPaths());
-    return pathValues;
   }
 
   @Override

entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DataJobInfoPatchBuilder.java

@@ -4,13 +4,11 @@ import static com.fasterxml.jackson.databind.node.JsonNodeFactory.instance;
 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 class DataJobInfoPatchBuilder extends AbstractMultiFieldPatchBuilder<Data
     return this;
   }
 
-  @Override
-  protected List<ImmutableTriple<String, String, JsonNode>> getPathValues() {
-    pathValues.addAll(customPropertiesPatchBuilder.getSubPaths());
-    return pathValues;
-  }
-
   @Override
   protected String getAspectName() {
     return DATA_JOB_INFO_ASPECT_NAME;

entity-registry/src/main/java/com/linkedin/metadata/aspect/patch/builder/DatasetPropertiesPatchBuilder.java

@@ -4,10 +4,8 @@ import static com.fasterxml.jackson.databind.node.JsonNodeFactory.instance;
 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 class DatasetPropertiesPatchBuilder
     return this;
   }
 
-  @Override
-  protected List<ImmutableTriple<String, String, JsonNode>> getPathValues() {
-    pathValues.addAll(customPropertiesPatchBuilder.getSubPaths());
-    return pathValues;
-  }
-
   @Override
   protected String getAspectName() {
     return DATASET_PROPERTIES_ASPECT_NAME;

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<T extends AbstractMultiFieldPatchBuild
 
   /** Convenience method to return parent patch builder in functional callstack */
   T getParent();
-
-  /**
-   * Exposes subpath values to parent patch builder in Op, Path, Value triples. Should usually only
-   * be called by the parent patch builder class when constructing the path values.
-   */
-  List<ImmutableTriple<String, String, JsonNode>> getSubPaths();
 }

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<ImmutableTriple<String, String, JsonNode>> 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<ImmutableTriple<String, String, JsonNode>> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> operation = pathValues.get(0);
+    assertEquals(operation.getLeft(), "remove");
+    assertEquals(operation.getMiddle(), "/customProperties/key1");
+    assertNull(operation.getRight());
+  }
+
+  @Test
+  public void testSetCustomProperties() {
+    Map<String, String> properties = new HashMap<>();
+    properties.put("key1", "value1");
+    properties.put("key2", "value2");
+
+    builder.setCustomProperties(properties);
+    builder.build();
+
+    List<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> operation = pathValues.get(0);
+    assertEquals(operation.getLeft(), "add");
+    assertEquals(operation.getMiddle(), "/customProperties");
+    assertTrue(operation.getRight().isObject());
+  }
+}

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<ImmutableTriple<String, String, JsonNode>> 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<ImmutableTriple<String, String, JsonNode>> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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<ImmutableTriple<String, String, JsonNode>> pathValues = builder.getTestPathValues();
+    assertNotNull(pathValues);
+    assertEquals(pathValues.size(), 1);
+
+    ImmutableTriple<String, String, JsonNode> 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);
+        });
+  }
+}