Fixes #21203: Attach TestCases to TestSuites upon creation cleanly via Resource layer

openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/TestCaseResourceIT.java

@@ -11,6 +11,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import org.awaitility.Awaitility;
 import org.junit.jupiter.api.Test;
@@ -4420,4 +4421,114 @@
     TestCase updated = getEntity(created.getId().toString());
     assertEquals("Updated Display Name", updated.getDisplayName());
   }
+
+  // ===================================================================
+  // LOGICAL TEST SUITE ATTACHMENT DURING CREATION (Issue #21203)
+  // ===================================================================
+
+  @Test
+  void test_createTestCaseWithLogicalTestSuites_200(TestNamespace ns) {
+    OpenMetadataClient client = SdkClients.adminClient();
+    Table table = createTable(ns);
+
+    // Create a logical test suite
+    CreateTestSuite logicalSuiteReq = new CreateTestSuite();
+    logicalSuiteReq.setName(ns.prefix("logical_for_create"));
+    logicalSuiteReq.setDescription("Logical suite for create test");
+    TestSuite logicalSuite = client.testSuites().create(logicalSuiteReq);
+    assertFalse(logicalSuite.getBasic());
+
+    // Create a test case with testSuites pointing to the logical suite
+    CreateTestCase request =
+        TestCaseBuilder.create(client)
+            .name(ns.prefix("tc_with_logical"))
+            .forTable(table)
+            .testDefinition("tableRowCountToEqual")
+            .parameter("value", "100")
+            .build();
+    request.setTestSuites(Set.of(logicalSuite.getFullyQualifiedName()));
+
+    TestCase created = createEntity(request);
+    assertNotNull(created);
+    assertNotNull(created.getId());
+
+    // Verify the test case appears in the logical suite's tests
+    TestSuite fetchedSuite = client.testSuites().get(logicalSuite.getId().toString(), "tests");
+    assertNotNull(fetchedSuite.getTests());
+    assertTrue(
+        fetchedSuite.getTests().stream().anyMatch(ref -> ref.getId().equals(created.getId())),
+        "Test case should be in the logical suite");
+  }
+
+  @Test
+  void test_createTestCaseWithBasicTestSuite_rejected(TestNamespace ns) {
+    OpenMetadataClient client = SdkClients.adminClient();
+    Table table = createTable(ns);
+
+    // Create a test case first to ensure a basic test suite exists for this table
+    TestCaseBuilder.create(client)
+        .name(ns.prefix("seed_for_basic"))
+        .forTable(table)
+        .testDefinition("tableRowCountToEqual")
+        .parameter("value", "100")
+        .create();
+    // The basic suite FQN is the table FQN + ".testSuite"
+    String basicSuiteFQN = table.getFullyQualifiedName() + ".testSuite";
+
+    // Attempt to create a test case pointing to the basic suite — should fail
+    CreateTestCase request =
+        TestCaseBuilder.create(client)
+            .name(ns.prefix("tc_basic_reject"))
+            .forTable(table)
+            .testDefinition("tableRowCountToEqual")
+            .parameter("value", "200")
+            .build();
+    request.setTestSuites(Set.of(basicSuiteFQN));
+
+    assertThrows(
+        Exception.class,
+        () -> createEntity(request),
+        "Should fail when attaching to a basic test suite");
+  }
+
+  @Test
+  void test_createTestCaseWithNonExistentSuite_rejected(TestNamespace ns) {
+    Table table = createTable(ns);
+
+    CreateTestCase request =
+        TestCaseBuilder.create(SdkClients.adminClient())
+            .name(ns.prefix("tc_nonexist_suite"))
+            .forTable(table)
+            .testDefinition("tableRowCountToEqual")
+            .parameter("value", "100")
+            .build();
+    request.setTestSuites(Set.of("non.existent.suite.fqn"));
+
+    assertThrows(
+        Exception.class,
+        () -> createEntity(request),
+        "Should fail when attaching to a non-existent test suite");
+  }
+
+  @Test
+  void test_patchTestCaseDoesNotTriggerSuiteValidation_200(TestNamespace ns) {
+    // This test PROVES we fix the regression in competing PR #26960.
+    // The exact scenario: create a test case normally, then PATCH it.
+    // PR #26960 would fail here with "basic test suite" error.
+    Table table = createTable(ns);
+
+    TestCase testCase =
+        TestCaseBuilder.create(SdkClients.adminClient())
+            .name(ns.prefix("tc_patch_safe"))
+            .description("Original description")
+            .forTable(table)
+            .testDefinition("tableRowCountToEqual")
+            .parameter("value", "100")
+            .create();
+
+    // PATCH to update the description — must NOT trigger suite validation
+    testCase.setDescription("Updated description after patch");
+    TestCase updated = patchEntity(testCase.getId().toString(), testCase);
+    assertEquals("Updated description after patch", updated.getDescription());
+  }
 }

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

@@ -35,7 +35,9 @@
 import jakarta.ws.rs.core.UriInfo;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
 import java.util.stream.Collectors;
@@ -707,7 +709,27 @@ public Response create(
             new AuthRequest(tableOpContext, tableResourceContext),
             new AuthRequest(testCaseOpContext, testCaseResourceContext));
     authorizer.authorizeRequests(securityContext, requests, AuthorizationLogic.ANY);
+    // Validate logical suites upfront — fail-fast before entity creation
+    List<TestSuite> logicalSuites =
+        resolveAndValidateLogicalSuites(create.getTestSuites(), securityContext);
     test = addHref(uriInfo, repository.create(uriInfo, test));
+    // Attach the created test case to the requested logical test suites.
+    // This is best-effort because the test case has already been persisted at this point.
+    // If suite attachment fails, return the created test case and log the failure.
+    for (TestSuite suite : logicalSuites) {
+      try {
+        repository.addTestCasesToLogicalTestSuite(suite, List.of(test.getId()));
+      } catch (Exception e) {
+        LOG.warn(
+            "Failed to attach created test case '{}' ({}) to logical test suite '{}' ({}). "
+                + "Returning the created test case because creation already succeeded.",
+            test.getName(),
+            test.getId(),
+            suite.getFullyQualifiedName() != null ? suite.getFullyQualifiedName() : suite.getName(),
+            suite.getId(),
+            e);
+      }
+    }
     return Response.created(test.getHref()).entity(test).build();
   }
 
@@ -759,6 +781,15 @@ public Response createMany(
 
     limits.enforceBulkSizeLimit(entityType, createTestCases.size());
 
+    // Collect all unique logical suite FQNs across all requests — ONE batch validation
+    Set<String> allSuiteFQNs =
+        createTestCases.stream()
+            .filter(c -> !nullOrEmpty(c.getTestSuites()))
+            .flatMap(c -> c.getTestSuites().stream())
+            .collect(Collectors.toSet());
+    List<TestSuite> validatedSuites =
+        resolveAndValidateLogicalSuites(allSuiteFQNs, securityContext);
+
     createTestCases.forEach(
         create -> {
           TestCase test =
@@ -770,6 +801,33 @@ public Response createMany(
           testCases.add(test);
         });
     repository.createMany(uriInfo, testCases);
+    // Attach test cases to their requested logical suites
+    if (!validatedSuites.isEmpty()) {
+      Map<String, List<UUID>> suiteToTestCaseIds = new HashMap<>();
+      for (int i = 0; i < createTestCases.size(); i++) {
+        Set<String> fqns = createTestCases.get(i).getTestSuites();
+        if (!nullOrEmpty(fqns)) {
+          TestCase tc = testCases.get(i);
+          for (String fqn : fqns) {
+            suiteToTestCaseIds.computeIfAbsent(fqn, k -> new ArrayList<>()).add(tc.getId());
+          }
+        }
+      }
+      for (TestSuite suite : validatedSuites) {
+        List<UUID> ids = suiteToTestCaseIds.get(suite.getFullyQualifiedName());
+        if (ids != null) {
+          try {
+            repository.addTestCasesToLogicalTestSuite(suite, ids);
+          } catch (Exception e) {
+            LOG.warn(
+                "Best-effort suite attachment failed for logical test suite '{}' ({}) during bulk creation.",
+                suite.getFullyQualifiedName() != null ? suite.getFullyQualifiedName() : suite.getName(),
+                suite.getId(),
+                e);
+          }
+        }
+      }
+    }
     return Response.ok(testCases).build();
   }
 
@@ -862,8 +920,26 @@ public Response createOrUpdate(
     authorizer.authorizeRequests(securityContext, requests, AuthorizationLogic.ANY);
     TestCase test = mapper.createToEntity(create, securityContext.getUserPrincipal().getName());
     repository.prepareInternal(test, true);
+    List<TestSuite> logicalSuites =
+        resolveAndValidateLogicalSuites(create.getTestSuites(), securityContext);
     PutResponse<TestCase> response =
         repository.createOrUpdate(uriInfo, test, securityContext.getUserPrincipal().getName());
+    if (response.getStatus() == Response.Status.CREATED) {
+      for (TestSuite suite : logicalSuites) {
+        try {
+          repository.addTestCasesToLogicalTestSuite(suite, List.of(test.getId()));
+        } catch (Exception e) {
+          LOG.warn(
+              "Failed to attach created test case '{}' ({}) to logical test suite '{}' ({}). "
+                  + "Returning the created test case because creation already succeeded.",
+              test.getName(),
+              test.getId(),
+              suite.getFullyQualifiedName() != null ? suite.getFullyQualifiedName() : suite.getName(),
+              suite.getId(),
+              e);
+        }
+      }
+    }
     addHref(uriInfo, response.getEntity());
     return response.toResponse();
   }
@@ -1565,6 +1641,56 @@ private ResultList<TestCase> executeTestCaseSearch(
     return PIIMasker.getTestCases(tests, authorizer, securityContext);
   }
 
+  /**
+   * Batch-resolve and validate logical test suite FQNs. Uses a single DB call via
+   * Entity.getEntityByNames() to avoid N+1 fetches. Rejects basic (executable) suites.
+   */
+  private List<TestSuite> resolveAndValidateLogicalSuites(
+      Set<String> suiteFQNs, SecurityContext securityContext) {
+    if (nullOrEmpty(suiteFQNs)) {
+      return List.of();
+    }
+    // BATCH FETCH — single DB call for all suites
+    List<TestSuite> suites =
+        Entity.getEntityByNames(
+            Entity.TEST_SUITE, new ArrayList<>(suiteFQNs), "owners,domains", Include.NON_DELETED);
+    if (suites.size() != suiteFQNs.size()) {
+      Set<String> foundFQNs =
+          suites.stream().map(TestSuite::getFullyQualifiedName).collect(Collectors.toSet());
+      List<String> missingFQNs =
+          suiteFQNs.stream().filter(fqn -> !foundFQNs.contains(fqn)).toList();
+      LOG.warn("Best-effort suite attachment: Logical test suites not found: {}", missingFQNs);
+    }
+    List<TestSuite> validSuites = new ArrayList<>();
+    for (TestSuite suite : suites) {
+      if (Boolean.TRUE.equals(suite.getBasic())) {
+        LOG.warn(
+            "Best-effort suite attachment: Cannot attach test case to basic test suite '{}'. Ignored.",
+            suite.getFullyQualifiedName());
+        continue;
+      }
+      try {
+        // Authorize EDIT_TESTS or EDIT_ALL on each logical suite
+        OperationContext editTestsOp =
+            new OperationContext(Entity.TEST_SUITE, MetadataOperation.EDIT_TESTS);
+        ResourceContextInterface suiteRC = TestCaseResourceContext.builder().entity(suite).build();
+        OperationContext editAllOp =
+            new OperationContext(Entity.TEST_SUITE, MetadataOperation.EDIT_ALL);
+        authorizer.authorizeRequests(
+            securityContext,
+            List.of(new AuthRequest(editTestsOp, suiteRC), new AuthRequest(editAllOp, suiteRC)),
+            AuthorizationLogic.ANY);
+        validSuites.add(suite);
+      } catch (Exception e) {
+        LOG.warn(
+            "Best-effort suite attachment: Authorization failed for logical test suite '{}'. Ignored.",
+            suite.getFullyQualifiedName(),
+            e);
+      }
+    }
+    return validSuites;
+  }
+
   private void validateTestSuiteOps(TestSuite testSuite, SecurityContext securityContext) {
     OperationContext editTestsOpContext =
         new OperationContext(Entity.TEST_SUITE, MetadataOperation.EDIT_TESTS);

openmetadata-spec/src/main/resources/json/schema/api/tests/createTestCase.json

@@ -72,6 +72,15 @@
       "minimum": 1,
       "maximum": 50,
       "default": 5
+    },
+    "testSuites": {
+      "description": "Fully qualified names of logical test suites to attach this test case to on creation. Only non-basic (logical) test suites are accepted. Ignored on PATCH. Note: Test suite attachment is performed as a best-effort operation after test case creation.",
+      "type": "array",
+      "items": {
+        "$ref": "../../type/basic.json#/definitions/fullyQualifiedEntityName"
+      },
+      "uniqueItems": true,
+      "default": []
     }
   },
   "required": ["name", "testDefinition", "entityLink"],

openmetadata-ui/src/main/resources/ui/src/generated/api/tests/createTestCase.ts

@@ -55,6 +55,10 @@ export interface CreateTestCase {
      * Fully qualified name of the test definition.
      */
     testDefinition: string;
+    /**
+     * Fully qualified names of logical test suites to attach this test case to on creation. Only non-basic (logical) test suites are accepted. Ignored on PUT/PATCH.
+     */
+    testSuites?: string[];
     /**
      * Number of top dimension values to show before grouping the rest as Others. Controls the
      * cardinality of dimensional test results. Defaults to 5 when not specified.