Split test class to solve flakiness (#859)

Problem:
The tests in this class were failing frequently, and it seemed to be due to test setup.
The test setup creates an Elasticsearch container, used by all the tests.
Except, one of the tests testSsl, so it re-creates the container with SSL enabled.

Solution:
Move testSsl to a different class for isolation.

Author: sharanrudresh

src/test/java/io/confluent/connect/elasticsearch/ElasticsearchClientSslTest.java

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2025 Confluent Inc.
+ *
+ * Licensed under the Confluent Community License (the "License"); you may not use
+ * this file except in compliance with the License.  You may obtain a copy of the
+ * License at
+ *
+ * http://www.confluent.io/confluent-community-license
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OF ANY KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+
+package io.confluent.connect.elasticsearch;
+
+import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.CONNECTION_PASSWORD_CONFIG;
+import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.CONNECTION_URL_CONFIG;
+import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.CONNECTION_USERNAME_CONFIG;
+import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.SECURITY_PROTOCOL_CONFIG;
+import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.SSL_CONFIG_PREFIX;
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.mock;
+
+import java.util.HashMap;
+
+import org.apache.kafka.common.config.SslConfigs;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.SecurityProtocol;
+import io.confluent.connect.elasticsearch.helper.ElasticsearchContainer;
+import io.confluent.connect.elasticsearch.helper.ElasticsearchHelperClient;
+
+public class ElasticsearchClientSslTest extends ElasticsearchClientTestBase {
+
+  private static final String ELASTIC_SUPERUSER_NAME = "elastic";
+  private static final String ELASTIC_SUPERUSER_PASSWORD = "elastic";
+
+  @BeforeClass
+  public static void setupBeforeAll() {
+    container = ElasticsearchContainer.fromSystemProperties().withSslEnabled(true);
+    container.start();
+  }
+
+  @Override
+  @Before
+  public void setup() {
+    index = TOPIC;
+    props = ElasticsearchSinkConnectorConfigTest.addNecessaryProps(new HashMap<>());
+    String address = container.getConnectionUrl(false);
+    props.put(CONNECTION_URL_CONFIG, address);
+    props.put(CONNECTION_USERNAME_CONFIG, ELASTIC_SUPERUSER_NAME);
+    props.put(CONNECTION_PASSWORD_CONFIG, ELASTIC_SUPERUSER_PASSWORD);
+    props.put(SECURITY_PROTOCOL_CONFIG, SecurityProtocol.SSL.name());
+    props.put(SSL_CONFIG_PREFIX + SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, container.getKeystorePath());
+    props.put(SSL_CONFIG_PREFIX + SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG, container.getKeystorePassword());
+    props.put(SSL_CONFIG_PREFIX + SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG, container.getTruststorePath());
+    props.put(SSL_CONFIG_PREFIX + SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG, container.getTruststorePassword());
+    props.put(SSL_CONFIG_PREFIX + SslConfigs.SSL_KEY_PASSWORD_CONFIG, container.getKeyPassword());
+    config = new ElasticsearchSinkConnectorConfig(props);
+    converter = new DataConverter(config);
+    helperClient = new ElasticsearchHelperClient(address, config,
+        container.shouldStartClientInCompatibilityMode());
+    helperClient.waitForConnection(30000);
+    offsetTracker = mock(OffsetTracker.class);
+  }
+
+  @Test
+  public void testSsl() throws Exception {
+    ElasticsearchClient client = new ElasticsearchClient(config, null, 
+    () -> offsetTracker.updateOffsets());
+    client.createIndexOrDataStream(index);
+
+    writeRecord(sinkRecord(0), client);
+    client.flush();
+
+    waitUntilRecordsInES(1);
+    assertEquals(1, helperClient.getDocCount(index));
+    client.close();
+  }
+} 

src/test/java/io/confluent/connect/elasticsearch/ElasticsearchClientTest.java

@@ -18,9 +18,7 @@
 import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.BATCH_SIZE_CONFIG;
 import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.BEHAVIOR_ON_MALFORMED_DOCS_CONFIG;
 import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.BEHAVIOR_ON_NULL_VALUES_CONFIG;
-import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.CONNECTION_PASSWORD_CONFIG;
 import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.CONNECTION_URL_CONFIG;
-import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.CONNECTION_USERNAME_CONFIG;
 import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.DATA_STREAM_DATASET_CONFIG;
 import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.DATA_STREAM_TYPE_CONFIG;
 import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.IGNORE_KEY_CONFIG;
@@ -29,8 +27,6 @@
 import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.MAX_IN_FLIGHT_REQUESTS_CONFIG;
 import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.MAX_RETRIES_CONFIG;
 import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.RETRY_BACKOFF_MS_CONFIG;
-import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.SECURITY_PROTOCOL_CONFIG;
-import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.SSL_CONFIG_PREFIX;
 import static io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.WRITE_METHOD_CONFIG;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -46,7 +42,6 @@
 
 import io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.BehaviorOnMalformedDoc;
 import io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.BehaviorOnNullValues;
-import io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.SecurityProtocol;
 import io.confluent.connect.elasticsearch.ElasticsearchSinkConnectorConfig.WriteMethod;
 import io.confluent.connect.elasticsearch.helper.ElasticsearchContainer;
 import io.confluent.connect.elasticsearch.helper.ElasticsearchHelperClient;
@@ -58,55 +53,29 @@
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
-import org.apache.kafka.common.config.SslConfigs;
-import org.apache.kafka.common.record.TimestampType;
 import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.data.SchemaBuilder;
 import org.apache.kafka.connect.data.Struct;
 import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.sink.ErrantRecordReporter;
 import org.apache.kafka.connect.sink.SinkRecord;
-import org.apache.kafka.test.TestUtils;
-import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.DocWriteRequest;
 import org.elasticsearch.action.bulk.BulkItemResponse;
 import org.elasticsearch.index.VersionType;
 import org.elasticsearch.search.SearchHit;
 import org.junit.After;
-import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 
-public class ElasticsearchClientTest {
-
-  private static final String INDEX = "index";
-  private static final String ELASTIC_SUPERUSER_NAME = "elastic";
-  private static final String ELASTIC_SUPERUSER_PASSWORD = "elastic";
-  private static final String TOPIC = "index";
-  private static final String DATA_STREAM_TYPE = "logs";
-  private static final String DATA_STREAM_DATASET = "dataset";
-
-  private static ElasticsearchContainer container;
-
-  private DataConverter converter;
-  private ElasticsearchHelperClient helperClient;
-  private ElasticsearchSinkConnectorConfig config;
-  private Map<String, String> props;
-  private String index;
-  private OffsetTracker offsetTracker;
+public class ElasticsearchClientTest extends ElasticsearchClientTestBase {
 
   @BeforeClass
   public static void setupBeforeAll() {
     container = ElasticsearchContainer.fromSystemProperties();
     container.start();
   }
 
-  @AfterClass
-  public static void cleanupAfterAll() {
-    container.close();
-  }
-
   @Before
   public void setup() {
     index = TOPIC;
@@ -747,43 +716,6 @@ public void testNoVersionConflict() throws Exception {
     client2.close();
   }
 
-  @Test
-  public void testSsl() throws Exception {
-    container.close();
-    container = ElasticsearchContainer.fromSystemProperties().withSslEnabled(true);
-    container.start();
-
-    String address = container.getConnectionUrl(false);
-    props.put(CONNECTION_URL_CONFIG, address);
-    props.put(CONNECTION_USERNAME_CONFIG, ELASTIC_SUPERUSER_NAME);
-    props.put(CONNECTION_PASSWORD_CONFIG, ELASTIC_SUPERUSER_PASSWORD);
-    props.put(SECURITY_PROTOCOL_CONFIG, SecurityProtocol.SSL.name());
-    props.put(SSL_CONFIG_PREFIX + SslConfigs.SSL_KEYSTORE_LOCATION_CONFIG, container.getKeystorePath());
-    props.put(SSL_CONFIG_PREFIX + SslConfigs.SSL_KEYSTORE_PASSWORD_CONFIG, container.getKeystorePassword());
-    props.put(SSL_CONFIG_PREFIX + SslConfigs.SSL_TRUSTSTORE_LOCATION_CONFIG, container.getTruststorePath());
-    props.put(SSL_CONFIG_PREFIX + SslConfigs.SSL_TRUSTSTORE_PASSWORD_CONFIG, container.getTruststorePassword());
-    props.put(SSL_CONFIG_PREFIX + SslConfigs.SSL_KEY_PASSWORD_CONFIG, container.getKeyPassword());
-    config = new ElasticsearchSinkConnectorConfig(props);
-    converter = new DataConverter(config);
-
-    ElasticsearchClient client = new ElasticsearchClient(config, null, () -> offsetTracker.updateOffsets());
-    helperClient = new ElasticsearchHelperClient(address, config,
-        container.shouldStartClientInCompatibilityMode());
-    client.createIndexOrDataStream(index);
-
-    writeRecord(sinkRecord(0), client);
-    client.flush();
-
-    waitUntilRecordsInES(1);
-    assertEquals(1, helperClient.getDocCount(index));
-    client.close();
-    helperClient = null;
-
-    container.close();
-    container = ElasticsearchContainer.fromSystemProperties();
-    container.start();
-  }
-
   @Test
   public void testWriteDataStreamInjectTimestamp() throws Exception {
     props.put(DATA_STREAM_TYPE_CONFIG, DATA_STREAM_TYPE);
@@ -805,72 +737,11 @@ public void testWriteDataStreamInjectTimestamp() throws Exception {
     client.close();
   }
 
-  private String createIndexName(String name) {
-    return config.isDataStream()
-        ? String.format("%s-%s-%s", DATA_STREAM_TYPE, DATA_STREAM_DATASET, name)
-        : name;
-  }
-
   @Test
   public void testConnectionUrlExtraSlash() {
     props.put(CONNECTION_URL_CONFIG, container.getConnectionUrl() + "/");
     config = new ElasticsearchSinkConnectorConfig(props);
     ElasticsearchClient client = new ElasticsearchClient(config, null, () -> offsetTracker.updateOffsets());
     client.close();
   }
-
-  private static Schema schema() {
-    return SchemaBuilder
-        .struct()
-        .name("record")
-        .field("offset", SchemaBuilder.int32().defaultValue(0).build())
-        .field("another", SchemaBuilder.int32().defaultValue(0).build())
-        .build();
-  }
-
-  private static SinkRecord sinkRecord(int offset) {
-    return sinkRecord("key", offset);
-  }
-
-  private static SinkRecord sinkRecord(String key, int offset) {
-    Struct value = new Struct(schema()).put("offset", offset).put("another", offset + 1);
-    return sinkRecord(key, schema(), value, offset);
-  }
-
-  private static SinkRecord sinkRecord(String key, Schema schema, Struct value, int offset) {
-    return new SinkRecord(
-        TOPIC,
-        0,
-        Schema.STRING_SCHEMA,
-        key,
-        schema,
-        value,
-        offset,
-        System.currentTimeMillis(),
-        TimestampType.CREATE_TIME
-    );
-  }
-
-  private void waitUntilRecordsInES(int expectedRecords) throws InterruptedException {
-    TestUtils.waitForCondition(
-        () -> {
-          try {
-            return helperClient.getDocCount(index) == expectedRecords;
-          } catch (ElasticsearchStatusException e) {
-            if (e.getMessage().contains("index_not_found_exception")) {
-              return false;
-            }
-
-            throw e;
-          }
-        },
-        TimeUnit.MINUTES.toMillis(1),
-        String.format("Could not find expected documents (%d) in time.", expectedRecords)
-    );
-  }
-
-  private void writeRecord(SinkRecord record, ElasticsearchClient client) {
-    client.index(record, converter.convertRecord(record, createIndexName(record.topic())),
-            new AsyncOffsetTracker.AsyncOffsetState(record.kafkaOffset()));
-  }
 }