Fixes #21203: Add testSuites field to CreateTestCaseRequest API for logical s…

.github/workflows/maven-build-collate.yml

@@ -79,11 +79,12 @@ jobs:
         uses: the-actions-org/workflow-dispatch@v4
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-          SHA: ${{ github.event_name == 'push' && github.event.after || github.event.pull_request.head.sha }}
+          SHA: ${{ github.event_name == 'push' && github.event.after || (github.event.pull_request.head.repo.full_name == github.repository && github.event.pull_request.head.sha || github.event.pull_request.merge_commit_sha || github.event.pull_request.base.sha) }}
+          COLLATE_EVENT: ${{ github.event_name == 'push' && 'push' || 'pull_request' }}
         with:
           workflow: OpenMetadata Collate Test
           ref: main
           repo: open-metadata/openmetadata-collate
           token: ${{ secrets.COLLATE_PAT }}
           wait-for-completion: true
-          inputs: '{ "sha": "${{ env.SHA }}", "event": "${{ github.event_name }}" }'
+          inputs: '{ "sha": "${{ env.SHA }}", "event": "${{ env.COLLATE_EVENT }}" }'

ingestion/Dockerfile

@@ -60,9 +60,15 @@ ENV LD_LIBRARY_PATH=/instantclient
 # Install DB2 iAccess Driver
 RUN if [ $(uname -m) = "x86_64" ]; \
   then \
-  curl https://public.dhe.ibm.com/software/ibmi/products/odbc/debs/dists/1.1.0/ibmi-acs-1.1.0.list | tee /etc/apt/sources.list.d/ibmi-acs-1.1.0.list \
-  && apt update \
-  && apt install ibm-iaccess; \
+  if curl -f --connect-timeout 10 --max-time 30 \
+    https://public.dhe.ibm.com/software/ibmi/products/odbc/debs/dists/1.1.0/ibmi-acs-1.1.0.list \
+    -o /etc/apt/sources.list.d/ibmi-acs-1.1.0.list; then \
+    apt update \
+    && apt install -y ibm-iaccess \
+    || echo "Warning: Failed to install ibm-iaccess, continuing without it"; \
+  else \
+    echo "Warning: Failed to fetch IBM iAccess repo list, continuing without ibm-iaccess"; \
+  fi; \
   fi
 
 # Required for Starting Ingestion Container in Docker Compose

ingestion/Dockerfile.ci

@@ -60,9 +60,15 @@ ENV LD_LIBRARY_PATH=/instantclient
 # Install DB2 iAccess Driver
 RUN if [ $(uname -m) = "x86_64" ]; \
   then \
-  curl https://public.dhe.ibm.com/software/ibmi/products/odbc/debs/dists/1.1.0/ibmi-acs-1.1.0.list | tee /etc/apt/sources.list.d/ibmi-acs-1.1.0.list \
-  && apt update \
-  && apt install ibm-iaccess; \
+  if curl -f --connect-timeout 10 --max-time 30 \
+    https://public.dhe.ibm.com/software/ibmi/products/odbc/debs/dists/1.1.0/ibmi-acs-1.1.0.list \
+    -o /etc/apt/sources.list.d/ibmi-acs-1.1.0.list; then \
+    apt update \
+    && apt install -y ibm-iaccess \
+    || echo "Warning: Failed to install ibm-iaccess, continuing without it"; \
+  else \
+    echo "Warning: Failed to fetch IBM iAccess repo list, continuing without ibm-iaccess"; \
+  fi; \
   fi
 
 # Required for Starting Ingestion Container in Docker Compose

ingestion/operators/docker/Dockerfile

@@ -63,16 +63,18 @@ RUN if [ $(uname -m) = "arm64" || $(uname -m) = "aarch64" ]; \
 ENV LD_LIBRARY_PATH=/instantclient
 
 # Install DB2 iAccess driver
-# The IBM repository approach is unreliable in CI environments, so we download the package directly
-# Use dpkg --force-depends because the package declares old Debian package names (libodbc1, odbcinst1debian2)
-# that don't exist in Debian 12, but the actual dependencies (unixodbc, odbcinst) are already installed
-RUN if [ $(uname -m) = "x86_64" ]; then \
-  wget -q https://public.dhe.ibm.com/software/ibmi/products/odbc/debs/dists/1.1.0/main/binary-amd64/ibm-iaccess-1.1.0.13-1.0.amd64.deb \
-    -O /tmp/ibm-iaccess.deb && \
-  dpkg -i --force-depends /tmp/ibm-iaccess.deb && \
-  apt-get install -f -y --no-install-recommends && \
-  rm -f /tmp/ibm-iaccess.deb; \
-fi
+RUN if [ $(uname -m) = "x86_64" ]; \
+  then \
+  if curl -f --connect-timeout 10 --max-time 30 \
+    https://public.dhe.ibm.com/software/ibmi/products/odbc/debs/dists/1.1.0/ibmi-acs-1.1.0.list \
+    -o /etc/apt/sources.list.d/ibmi-acs-1.1.0.list; then \
+    apt update \
+    && apt install -y ibm-iaccess \
+    || echo "Warning: Failed to install ibm-iaccess, continuing without it"; \
+  else \
+    echo "Warning: Failed to fetch IBM iAccess repo list, continuing without ibm-iaccess"; \
+  fi; \
+  fi
 
 WORKDIR ingestion/
 

ingestion/operators/docker/Dockerfile.ci

@@ -63,16 +63,18 @@ RUN if [ $(uname -m) = "arm64" ] | [ $(uname -m) = "aarch64" ]; \
 ENV LD_LIBRARY_PATH=/instantclient
 
 # Install DB2 iAccess Driver
-# The IBM repository approach is unreliable in CI environments, so we download the package directly
-# Use dpkg --force-depends because the package declares old Debian package names (libodbc1, odbcinst1debian2)
-# that don't exist in Debian 12, but the actual dependencies (unixodbc, odbcinst) are already installed
-RUN if [ $(uname -m) = "x86_64" ]; then \
-  wget -q https://public.dhe.ibm.com/software/ibmi/products/odbc/debs/dists/1.1.0/main/binary-amd64/ibm-iaccess-1.1.0.13-1.0.amd64.deb \
-    -O /tmp/ibm-iaccess.deb && \
-  dpkg -i --force-depends /tmp/ibm-iaccess.deb && \
-  apt-get install -f -y --no-install-recommends && \
-  rm -f /tmp/ibm-iaccess.deb; \
-fi
+RUN if [ $(uname -m) = "x86_64" ]; \
+  then \
+  if curl -f --connect-timeout 10 --max-time 30 \
+    https://public.dhe.ibm.com/software/ibmi/products/odbc/debs/dists/1.1.0/ibmi-acs-1.1.0.list \
+    -o /etc/apt/sources.list.d/ibmi-acs-1.1.0.list; then \
+    apt update \
+    && apt install -y ibm-iaccess \
+    || echo "Warning: Failed to install ibm-iaccess, continuing without it"; \
+  else \
+    echo "Warning: Failed to fetch IBM iAccess repo list, continuing without ibm-iaccess"; \
+  fi; \
+  fi
 
 WORKDIR /ingestion
 

ingestion/tests/integration/trino/conftest.py

@@ -206,7 +206,7 @@ def _execute_with_connect(sql):
     ).fetchall()
 
     _execute_with_connect(
-        "create schema minio.my_schema WITH (location = 's3a://hive-warehouse/')"
+        "create schema minio.my_schema WITH (location = 's3a://hive-warehouse/my_schema/')"
     )
     data_dir = os.path.dirname(__file__) + "/data"
     for file in os.listdir(data_dir):
@@ -218,7 +218,9 @@ def _execute_with_connect(sql):
             create_test_data_from_parquet(engine, file_path)
 
         sleep(1)
-        _execute_with_connect("ANALYZE " + f'minio."my_schema"."{file_path.stem}"')
+        retry(wait=wait_fixed(2), stop=stop_after_delay(120))(_execute_with_connect)(
+            "ANALYZE " + f'minio."my_schema"."{file_path.stem}"'
+        )
     _execute_with_connect(
         "CALL system.drop_stats(schema_name => 'my_schema', table_name => 'empty')"
     )

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

@@ -8,6 +8,7 @@
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
+import java.net.http.HttpTimeoutException;
 import java.time.Duration;
 import java.util.List;
 import java.util.UUID;
@@ -52,6 +53,8 @@ public class GlossaryOntologyExportIT {
   private static final Logger LOG = LoggerFactory.getLogger(GlossaryOntologyExportIT.class);
   private static final HttpClient HTTP_CLIENT =
       HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build();
+  private static final Duration DEFAULT_EXPORT_TIMEOUT = Duration.ofMinutes(3);
+  private static final Duration RDF_XML_EXPORT_TIMEOUT = Duration.ofMinutes(6);
 
   private static final String TURTLE_CONTENT_TYPE = "text/turtle";
   private static final String RDF_XML_CONTENT_TYPE = "application/rdf+xml";
@@ -428,11 +431,25 @@ private HttpResponse<String> exportGlossaryRaw(
             .uri(URI.create(url))
             .header("Authorization", "Bearer " + token)
             .header("Accept", acceptHeader)
-            .timeout(Duration.ofSeconds(60))
+            .timeout(getExportTimeout(format))
             .GET()
             .build();
 
-    return HTTP_CLIENT.send(request, HttpResponse.BodyHandlers.ofString());
+    return executeExportRequest(request);
+  }
+
+  private Duration getExportTimeout(String format) {
+    return "rdfxml".equalsIgnoreCase(format) || "xml".equalsIgnoreCase(format)
+        ? RDF_XML_EXPORT_TIMEOUT
+        : DEFAULT_EXPORT_TIMEOUT;
+  }
+
+  private HttpResponse<String> executeExportRequest(HttpRequest request) throws Exception {
+    try {
+      return HTTP_CLIENT.send(request, HttpResponse.BodyHandlers.ofString());
+    } catch (HttpTimeoutException timeoutException) {
+      return HTTP_CLIENT.send(request, HttpResponse.BodyHandlers.ofString());
+    }
   }
 
   private int countOccurrences(String text, String pattern) {

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

@@ -36,6 +36,7 @@
 import org.openmetadata.schema.type.Column;
 import org.openmetadata.schema.type.ColumnDataType;
 import org.openmetadata.schema.type.EntityHistory;
+import org.openmetadata.schema.type.EntityReference;
 import org.openmetadata.schema.type.TagLabel;
 import org.openmetadata.schema.type.csv.CsvImportResult;
 import org.openmetadata.schema.utils.JsonUtils;
@@ -45,6 +46,7 @@
 import org.openmetadata.sdk.models.ListResponse;
 import org.openmetadata.sdk.network.HttpMethod;
 import org.openmetadata.sdk.network.RequestOptions;
+import org.openmetadata.service.Entity;
 import org.openmetadata.service.resources.dqtests.TestCaseResource;
 
 /**
@@ -184,6 +186,24 @@ private Table createTable(TestNamespace ns, List<org.openmetadata.schema.type.Ta
     return SdkClients.adminClient().tables().create(tableRequest);
   }
 
+  private TestSuite createBasicTestSuite(CreateTestSuite request) {
+    OpenMetadataClient client = SdkClients.adminClient();
+
+    org.openmetadata.sdk.network.RequestOptions options =
+        org.openmetadata.sdk.network.RequestOptions.builder()
+            .queryParam("name", request.getName())
+            .build();
+
+    return client
+        .getHttpClient()
+        .execute(
+            org.openmetadata.sdk.network.HttpMethod.POST,
+            "/v1/dataQuality/testSuites/basic",
+            request,
+            TestSuite.class,
+            options);
+  }
+
   @Override
   protected TestCase createEntity(CreateTestCase createRequest) {
     return SdkClients.adminClient().testCases().create(createRequest);
@@ -954,6 +974,130 @@ void test_listTestCasesWithPagination(TestNamespace ns) {
     }
   }
 
+  @Test
+  void test_createTestCaseWithLogicalTestSuites(TestNamespace ns) throws Exception {
+    OpenMetadataClient client = SdkClients.adminClient();
+    Table table = createTable(ns);
+
+    CreateTestSuite logicalSuiteReq = new CreateTestSuite();
+    logicalSuiteReq.setName(ns.prefix("logical_suite_create_request"));
+    TestSuite logicalSuite = client.testSuites().create(logicalSuiteReq);
+
+    EntityReference logicalSuiteRef =
+        new EntityReference()
+            .withId(logicalSuite.getId())
+            .withType(Entity.TEST_SUITE)
+            .withName(logicalSuite.getName());
+
+    CreateTestCase request =
+        TestCaseBuilder.create(client)
+            .name(ns.prefix("create_suite_request_tc"))
+            .forTable(table)
+            .testDefinition("tableRowCountToEqual")
+            .parameter("value", "100")
+            .testSuites(List.of(logicalSuiteRef))
+            .build();
+
+    TestCase created = client.testCases().create(request);
+
+    Awaitility.await("test case is attached to logical test suite")
+        .atMost(Duration.ofSeconds(30))
+        .pollInterval(Duration.ofSeconds(2))
+        .untilAsserted(
+            () -> {
+              TestCase fetched = client.testCases().get(created.getId().toString(), "testSuites");
+              assertNotNull(fetched.getTestSuites());
+              assertTrue(
+                  fetched.getTestSuites().stream()
+                      .anyMatch(ts -> ts.getId().equals(logicalSuite.getId())));
+            });
+
+    TestSuite updatedSuite = client.testSuites().get(logicalSuite.getId().toString(), "tests");
+    assertNotNull(updatedSuite.getTests());
+    assertTrue(
+        updatedSuite.getTests().stream().anyMatch(ref -> ref.getId().equals(created.getId())));
+  }
+
+  @Test
+  void test_upsertTestCaseWithLogicalTestSuites(TestNamespace ns) throws Exception {
+    OpenMetadataClient client = SdkClients.adminClient();
+    Table table = createTable(ns);
+
+    CreateTestSuite logicalSuiteReq = new CreateTestSuite();
+    logicalSuiteReq.setName(ns.prefix("logical_suite_upsert_request"));
+    TestSuite logicalSuite = client.testSuites().create(logicalSuiteReq);
+
+    EntityReference logicalSuiteRef =
+        new EntityReference()
+            .withId(logicalSuite.getId())
+            .withType(Entity.TEST_SUITE)
+            .withName(logicalSuite.getName());
+
+    CreateTestCase createRequest =
+        TestCaseBuilder.create(client)
+            .name(ns.prefix("upsert_suite_request_tc"))
+            .forTable(table)
+            .testDefinition("tableRowCountToEqual")
+            .parameter("value", "100")
+            .description("Initial")
+            .testSuites(List.of(logicalSuiteRef))
+            .build();
+
+    TestCase created = client.testCases().upsert(createRequest);
+    assertNotNull(created);
+
+    CreateTestCase updateRequest =
+        TestCaseBuilder.create(client)
+            .name(ns.prefix("upsert_suite_request_tc"))
+            .forTable(table)
+            .testDefinition("tableRowCountToEqual")
+            .parameter("value", "200")
+            .description("Updated")
+            .testSuites(List.of(logicalSuiteRef))
+            .build();
+
+    TestCase updated = client.testCases().upsert(updateRequest);
+    assertEquals(created.getId(), updated.getId());
+    assertEquals("Updated", updated.getDescription());
+
+    Awaitility.await("test case upsert should preserve logical test suite membership")
+        .atMost(Duration.ofSeconds(30))
+        .pollInterval(Duration.ofSeconds(2))
+        .untilAsserted(
+            () -> {
+              TestCase fetched = client.testCases().get(created.getId().toString(), "testSuites");
+              assertNotNull(fetched.getTestSuites());
+              assertTrue(
+                  fetched.getTestSuites().stream()
+                      .anyMatch(ts -> ts.getId().equals(logicalSuite.getId())));
+            });
+  }
+
+  @Test
+  void test_createTestCaseWithLogicalTestSuites_basicSuiteRejected(TestNamespace ns) {
+    OpenMetadataClient client = SdkClients.adminClient();
+    Table table = createTable(ns);
+
+    CreateTestSuite basicSuiteReq = new CreateTestSuite();
+    basicSuiteReq.setName(table.getFullyQualifiedName());
+    basicSuiteReq.setBasicEntityReference(table.getFullyQualifiedName());
+    TestSuite basicSuite = createBasicTestSuite(basicSuiteReq);
+
+    EntityReference basicSuiteRef =
+        new EntityReference().withId(basicSuite.getId()).withType(Entity.TEST_SUITE);
+
+    CreateTestCase request =
+        TestCaseBuilder.create(client)
+            .name(ns.prefix("create_suite_request_basic_rejected_tc"))
+            .forTable(table)
+            .testDefinition("tableRowCountToEqual")
+            .parameter("value", "100")
+            .testSuites(List.of(basicSuiteRef))
+            .build();
+
+    assertThrows(Exception.class, () -> client.testCases().create(request));
+  }
+
   @Test
   void test_bulkAddTestCasesToLogicalTestSuiteByIds(TestNamespace ns) throws Exception {
     OpenMetadataClient client = SdkClients.adminClient();
@@ -1058,7 +1202,7 @@ void test_bulkAddTestCasesToLogicalTestSuiteByIds_basicSuiteRejected(TestNamespa
     CreateTestSuite basicSuiteReq = new CreateTestSuite();
     basicSuiteReq.setName(table.getFullyQualifiedName());
     basicSuiteReq.setBasicEntityReference(table.getFullyQualifiedName());
-    TestSuite basicSuite = client.testSuites().create(basicSuiteReq);
+    TestSuite basicSuite = createBasicTestSuite(basicSuiteReq);
 
     Map<String, Object> request = new HashMap<>();
     request.put("testSuiteId", basicSuite.getId().toString());

openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/builders/TestCaseBuilder.java

@@ -165,6 +165,14 @@ public TestCaseBuilder useDynamicAssertion(boolean value) {
     return this;
   }
 
+  /**
+   * Set test suite references (logical suites) for this test case.
+   */
+  public TestCaseBuilder testSuites(List<EntityReference> testSuites) {
+    request.setTestSuites(testSuites);
+    return this;
+  }
+
   /**
    * Build the CreateTestCase request without executing it.
    */

openmetadata-sdk/src/test/java/org/openmetadata/sdk/entities/TestCaseMockTest.java

@@ -48,7 +48,7 @@ void testCreateTestCase() {
     createRequest.setName("null-check-test");
     createRequest.setDisplayName("Null Check Test");
     createRequest.setDescription("Test to check for null values in column");
-    // createRequest.setTestSuite("quality-test-suite"); // Not available in CreateTestCase
+    // createRequest.setTestSuites(...); // Optional on CreateTestCase
 
     TestCase expectedTestCase = new TestCase();
     expectedTestCase.setId(UUID.randomUUID());