MINOR - Handle removed tests from contract & Add Rule Engine Tests (#22755)

* MINOR - Add Rule Engine Tests

* format

* improve test suite handling

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DataContractRepository.java

@@ -304,21 +304,27 @@ public class DataContractRepository extends EntityRepository<DataContract> {
     return fieldNames;
   }
 
+  private String getTestSuiteName(DataContract dataContract) {
+    return dataContract.getName() + " - Data Contract Expectations";
+  }
+
   private TestSuite createOrUpdateDataContractTestSuite(DataContract dataContract) {
-    if (nullOrEmpty(dataContract.getQualityExpectations())) {
-      return null; // No quality expectations, no test suite needed
-    }
     try {
+      // If we don't have quality expectations or a test suite, we don't need to create one
+      if (nullOrEmpty(dataContract.getQualityExpectations())
+          && !contractHasTestSuite(dataContract)) {
+        return null;
+      }
+
+      // If we had a test suite from older tests, but we removed them, we can delete the suite
+      if (nullOrEmpty(dataContract.getQualityExpectations())) {
+        deleteTestSuite(dataContract);
+        dataContract.setTestSuite(null);
+        return null;
+      }
 
-      TestCaseRepository testCaseRepository =
-          (TestCaseRepository) Entity.getEntityRepository(Entity.TEST_CASE);
       TestSuite testSuite = getOrCreateTestSuite(dataContract);
-
+      updateTestSuiteTests(dataContract, testSuite);
-      // Collect test case references from quality expectations
-      List<UUID> testCaseRefs =
-          dataContract.getQualityExpectations().stream().map(EntityReference::getId).toList();
-
-      testCaseRepository.addTestCasesToLogicalTestSuite(testSuite, testCaseRefs);
 
       // Add the test suite to the data contract
       dataContract.setTestSuite(
@@ -335,8 +341,45 @@ public class DataContractRepository extends EntityRepository<DataContract> {
     }
   }
 
+  private void updateTestSuiteTests(DataContract dataContract, TestSuite testSuite) {
+    TestCaseRepository testCaseRepository =
+        (TestCaseRepository) Entity.getEntityRepository(Entity.TEST_CASE);
+
+    // Collect test case references from quality expectations
+    List<UUID> testCaseRefs =
+        dataContract.getQualityExpectations().stream().map(EntityReference::getId).toList();
+    List<UUID> currentTests =
+        testSuite.getTests() != null
+            ? testSuite.getTests().stream().map(EntityReference::getId).toList()
+            : Collections.emptyList();
+
+    // Add only new tests to the test suite
+    List<UUID> newTestCases =
+        testCaseRefs.stream().filter(testCaseRef -> !currentTests.contains(testCaseRef)).toList();
+    if (!nullOrEmpty(newTestCases)) {
+      testCaseRepository.addTestCasesToLogicalTestSuite(testSuite, newTestCases);
+    }
+
+    // Then, remove any tests that are no longer in the quality expectations
+    List<UUID> testsToRemove =
+        currentTests.stream().filter(testId -> !testCaseRefs.contains(testId)).toList();
+    if (!nullOrEmpty(testsToRemove)) {
+      testsToRemove.forEach(
+          test -> {
+            testCaseRepository.deleteTestCaseFromLogicalTestSuite(testSuite.getId(), test);
+          });
+    }
+  }
+
+  private void deleteTestSuite(DataContract dataContract) {
+    TestSuiteRepository testSuiteRepository =
+        (TestSuiteRepository) Entity.getEntityRepository(Entity.TEST_SUITE);
+    TestSuite testSuite = getOrCreateTestSuite(dataContract);
+    testSuiteRepository.deleteLogicalTestSuite(ADMIN_USER_NAME, testSuite, true);
+  }
+
   private TestSuite getOrCreateTestSuite(DataContract dataContract) {
-    String testSuiteName = dataContract.getName() + " - Data Contract Expectations";
+    String testSuiteName = getTestSuiteName(dataContract);
     TestSuiteRepository testSuiteRepository =
         (TestSuiteRepository) Entity.getEntityRepository(Entity.TEST_SUITE);
 
@@ -371,6 +414,20 @@ public class DataContractRepository extends EntityRepository<DataContract> {
     return maybeTestSuite.get();
   }
 
+  private Boolean contractHasTestSuite(DataContract dataContract) {
+    TestSuiteRepository testSuiteRepository =
+        (TestSuiteRepository) Entity.getEntityRepository(Entity.TEST_SUITE);
+
+    String testSuiteName = getTestSuiteName(dataContract);
+
+    // Check if test suite already exists
+    Optional<TestSuite> maybeTestSuite =
+        testSuiteRepository.getByNameOrNull(
+            null, testSuiteName, testSuiteRepository.getFields(""), Include.NON_DELETED, false);
+
+    return maybeTestSuite.isPresent();
+  }
+
   // Prepare the Ingestion Pipeline from the test suite that will handle the execution
   private IngestionPipeline createIngestionPipeline(TestSuite testSuite) {
     IngestionPipelineRepository pipelineRepository =

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestSuiteRepository.java

@@ -527,10 +527,9 @@ public class TestSuiteRepository extends EntityRepository<TestSuite> {
   }
 
   public RestUtil.DeleteResponse<TestSuite> deleteLogicalTestSuite(
-      SecurityContext securityContext, TestSuite original, boolean hardDelete) {
+      String updatedBy, TestSuite original, boolean hardDelete) {
     // deleting a logical will delete the test suite and only remove the relationship to
     // test cases if hardDelete is true. Test Cases will not be deleted.
-    String updatedBy = securityContext.getUserPrincipal().getName();
     preDelete(original, updatedBy);
     setFieldsInternal(original, putFields);
     deleteChildIngestionPipelines(original.getId(), hardDelete, updatedBy);
@@ -581,7 +580,8 @@ public class TestSuiteRepository extends EntityRepository<TestSuite> {
         () -> {
           try {
             RestUtil.DeleteResponse<TestSuite> deleteResponse =
-                deleteLogicalTestSuite(securityContext, testSuite, hardDelete);
+                deleteLogicalTestSuite(
+                    securityContext.getUserPrincipal().getName(), testSuite, hardDelete);
             deleteFromSearch(deleteResponse.entity(), hardDelete);
 
             WebsocketNotificationHandler.sendDeleteOperationCompleteNotification(

openmetadata-service/src/main/java/org/openmetadata/service/resources/dqtests/TestSuiteResource.java

@@ -742,7 +742,8 @@ public class TestSuiteResource extends EntityResource<TestSuite, TestSuiteReposi
       throw new IllegalArgumentException(NON_BASIC_TEST_SUITE_DELETION_ERROR);
     }
     RestUtil.DeleteResponse<TestSuite> response =
-        repository.deleteLogicalTestSuite(securityContext, testSuite, hardDelete);
+        repository.deleteLogicalTestSuite(
+            securityContext.getUserPrincipal().getName(), testSuite, hardDelete);
     repository.deleteFromSearch(response.entity(), hardDelete);
     addHref(uriInfo, response.entity());
     return response.toResponse();
@@ -810,7 +811,8 @@ public class TestSuiteResource extends EntityResource<TestSuite, TestSuiteReposi
       throw new IllegalArgumentException(NON_BASIC_TEST_SUITE_DELETION_ERROR);
     }
     RestUtil.DeleteResponse<TestSuite> response =
-        repository.deleteLogicalTestSuite(securityContext, testSuite, hardDelete);
+        repository.deleteLogicalTestSuite(
+            securityContext.getUserPrincipal().getName(), testSuite, hardDelete);
     addHref(uriInfo, response.entity());
     return response.toResponse();
   }

openmetadata-service/src/test/java/org/openmetadata/service/resources/data/DataContractResourceTest.java

@@ -15,6 +15,7 @@ package org.openmetadata.service.resources.data;
 
 import static jakarta.ws.rs.core.Response.Status.BAD_REQUEST;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -41,6 +42,7 @@ import jakarta.ws.rs.core.Response.Status;
 import java.io.IOException;
 import java.net.URISyntaxException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
@@ -112,6 +114,7 @@ public class DataContractResourceTest extends EntityResourceTest<DataContract, C
 
   private static TestCaseResourceTest testCaseResourceTest;
   private static IngestionPipelineResourceTest ingestionPipelineResourceTest;
+  private static TestSuiteResourceTest testSuiteResourceTest;
   private static DataContractRepository dataContractRepository;
   private static PipelineServiceClientInterface mockPipelineClient;
 
@@ -130,7 +133,7 @@ public class DataContractResourceTest extends EntityResourceTest<DataContract, C
   @BeforeAll
   public void setup(TestInfo test) throws URISyntaxException, IOException {
     testCaseResourceTest = new TestCaseResourceTest();
-    // testCaseResourceTest.setup(test);
+    testSuiteResourceTest = new TestSuiteResourceTest();
     ingestionPipelineResourceTest = new IngestionPipelineResourceTest();
     ingestionPipelineResourceTest.setup(test);
 
@@ -950,6 +953,143 @@ public class DataContractResourceTest extends EntityResourceTest<DataContract, C
     assertNotNull(finalRetrieved.getTestSuite());
   }
 
+  @Test
+  @Execution(ExecutionMode.CONCURRENT)
+  void testUpdateTestSuiteTestsWithQualityExpectations(TestInfo test) throws IOException {
+    Table table = createUniqueTable(test.getDisplayName());
+    CreateDataContract create = createDataContractRequest(test.getDisplayName(), table);
+    DataContract created = createDataContract(create);
+
+    // Create initial test cases for quality expectations
+    String tableLink = String.format("<#E::table::%s>", table.getFullyQualifiedName());
+
+    CreateTestCase createTestCase1 =
+        testCaseResourceTest
+            .createRequest("test_case_completeness_" + test.getDisplayName())
+            .withEntityLink(tableLink);
+    TestCase testCase1 =
+        testCaseResourceTest.createAndCheckEntity(createTestCase1, ADMIN_AUTH_HEADERS);
+
+    CreateTestCase createTestCase2 =
+        testCaseResourceTest
+            .createRequest("test_case_validity_" + test.getDisplayName())
+            .withEntityLink(tableLink);
+    TestCase testCase2 =
+        testCaseResourceTest.createAndCheckEntity(createTestCase2, ADMIN_AUTH_HEADERS);
+
+    // Step 1: Add initial quality expectations with both test cases
+    List<EntityReference> initialQualityExpectations =
+        List.of(testCase1.getEntityReference(), testCase2.getEntityReference());
+
+    create.withStatus(ContractStatus.Active).withQualityExpectations(initialQualityExpectations);
+    DataContract updated = updateDataContract(create);
+
+    // Verify initial state - both test cases should be in TestSuite
+    assertEquals(2, updated.getQualityExpectations().size());
+    assertNotNull(updated.getTestSuite());
+
+    // Get TestSuite and verify it contains both tests
+    TestSuite testSuite =
+        testSuiteResourceTest.getEntity(
+            updated.getTestSuite().getId(), "tests", ADMIN_AUTH_HEADERS);
+
+    assertNotNull(testSuite.getTests());
+    assertEquals(2, testSuite.getTests().size());
+    assertTrue(testSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase1.getId())));
+    assertTrue(testSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase2.getId())));
+
+    // Step 2: Create a third test case
+    CreateTestCase createTestCase3 =
+        testCaseResourceTest
+            .createRequest("test_case_accuracy_" + test.getDisplayName())
+            .withEntityLink(tableLink);
+    TestCase testCase3 =
+        testCaseResourceTest.createAndCheckEntity(createTestCase3, ADMIN_AUTH_HEADERS);
+
+    // Update to add new test case while keeping existing ones
+    List<EntityReference> expandedQualityExpectations =
+        List.of(
+            testCase1.getEntityReference(),
+            testCase2.getEntityReference(),
+            testCase3.getEntityReference());
+
+    create.withQualityExpectations(expandedQualityExpectations);
+    DataContract updatedWithNewTest = updateDataContract(create);
+
+    // Verify new test case was added
+    assertEquals(3, updatedWithNewTest.getQualityExpectations().size());
+
+    // Verify TestSuite now contains all three tests
+    TestSuite expandedTestSuite =
+        testSuiteResourceTest.getEntity(
+            updated.getTestSuite().getId(), "tests", ADMIN_AUTH_HEADERS);
+
+    assertNotNull(expandedTestSuite.getTests());
+    assertEquals(3, expandedTestSuite.getTests().size());
+    assertTrue(
+        expandedTestSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase1.getId())));
+    assertTrue(
+        expandedTestSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase2.getId())));
+    assertTrue(
+        expandedTestSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase3.getId())));
+
+    // Step 3: Remove one test case (remove testCase2)
+    List<EntityReference> reducedQualityExpectations =
+        List.of(testCase1.getEntityReference(), testCase3.getEntityReference());
+
+    create.withQualityExpectations(reducedQualityExpectations);
+    DataContract updatedWithRemovedTest = updateDataContract(create);
+
+    // Verify test case was removed from quality expectations
+    assertEquals(2, updatedWithRemovedTest.getQualityExpectations().size());
+    assertTrue(
+        updatedWithRemovedTest.getQualityExpectations().stream()
+            .anyMatch(ref -> ref.getId().equals(testCase1.getId())));
+    assertTrue(
+        updatedWithRemovedTest.getQualityExpectations().stream()
+            .anyMatch(ref -> ref.getId().equals(testCase3.getId())));
+    assertFalse(
+        updatedWithRemovedTest.getQualityExpectations().stream()
+            .anyMatch(ref -> ref.getId().equals(testCase2.getId())));
+
+    // Verify TestSuite also had the test removed
+    TestSuite reducedTestSuite =
+        testSuiteResourceTest.getEntity(
+            updated.getTestSuite().getId(), "tests", ADMIN_AUTH_HEADERS);
+
+    assertNotNull(reducedTestSuite.getTests());
+    assertEquals(2, reducedTestSuite.getTests().size());
+    assertTrue(
+        reducedTestSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase1.getId())));
+    assertFalse(
+        reducedTestSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase2.getId())));
+    assertTrue(
+        reducedTestSuite.getTests().stream().anyMatch(t -> t.getId().equals(testCase3.getId())));
+
+    // Step 4: Remove all test cases
+    create.withQualityExpectations(Collections.emptyList());
+    DataContract updatedWithNoTests = updateDataContract(create);
+
+    // Verify all test cases were removed
+    assertTrue(
+        updatedWithNoTests.getQualityExpectations() == null
+            || updatedWithNoTests.getQualityExpectations().isEmpty());
+
+    // Verify TestSuite has been deleted
+    assertResponseContains(
+        () ->
+            testSuiteResourceTest.getEntity(
+                updated.getTestSuite().getId(), "tests", ADMIN_AUTH_HEADERS),
+        Status.NOT_FOUND,
+        "testSuite instance for " + updated.getTestSuite().getId().toString() + " not found");
+
+    // GET the data contract and verify persistence
+    DataContract retrieved = getDataContract(created.getId(), "");
+    assertTrue(
+        retrieved.getQualityExpectations() == null || retrieved.getQualityExpectations().isEmpty());
+    assertNull(retrieved.getTestSuite()); // We deleted the TestSuite when no tests remain
+  }
+
   @Test
   @Execution(ExecutionMode.CONCURRENT)
   @Override

openmetadata-service/src/test/java/org/openmetadata/service/rules/RuleEngineTests.java

@@ -3,7 +3,9 @@ package org.openmetadata.service.rules;
 import static org.junit.Assert.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.openmetadata.service.resources.EntityResourceTest.PII_SENSITIVE_TAG_LABEL;
 import static org.openmetadata.service.resources.EntityResourceTest.TEAM11_REF;
+import static org.openmetadata.service.resources.EntityResourceTest.TIER1_TAG_LABEL;
 import static org.openmetadata.service.resources.EntityResourceTest.USER1_REF;
 import static org.openmetadata.service.resources.EntityResourceTest.USER2_REF;
 import static org.openmetadata.service.util.TestUtils.ADMIN_AUTH_HEADERS;
@@ -433,4 +435,32 @@ public class RuleEngineTests extends OpenMetadataApplicationTest {
     // Ignored the table, should not pass
     Assertions.assertFalse(RuleEngine.getInstance().shouldApplyRule(table, ruleWithIgnoredTable));
   }
+
+  @Test
+  @Execution(ExecutionMode.CONCURRENT)
+  void testPIISensitiveTagRule(TestInfo test) {
+    Table table = getMockTable(test).withTags(List.of());
+
+    SemanticsRule piiRule =
+        new SemanticsRule()
+            .withName("Table has non-PII.Sensitive tags")
+            .withDescription("Table is not allowed to have PII.Sensitive tags")
+            .withRule(
+                "{\"!\":{\"some\":[{\"var\":\"tags\"},{\"==\":[{\"var\":\"tagFQN\"},\"PII.Sensitive\"]}]}}");
+
+    RuleEngine.getInstance().evaluate(table, List.of(piiRule), false, false);
+
+    table.withTags(List.of(PII_SENSITIVE_TAG_LABEL));
+    assertThrows(
+        RuleValidationException.class,
+        () -> RuleEngine.getInstance().evaluate(table, List.of(piiRule), false, false));
+
+    table.withTags(List.of(PII_SENSITIVE_TAG_LABEL, TIER1_TAG_LABEL));
+    assertThrows(
+        RuleValidationException.class,
+        () -> RuleEngine.getInstance().evaluate(table, List.of(piiRule), false, false));
+
+    table.withTags(List.of(TIER1_TAG_LABEL));
+    RuleEngine.getInstance().evaluate(table, List.of(piiRule), false, false);
+  }
 }