AbstractElasticsearchTest.java

/*******************************************************************************
 * Copyright (c) 2025 Eclipse RDF4J contributors.
 *
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Distribution License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/org/documents/edl-v10.php.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 *******************************************************************************/
// Some portions generated by Codex
package org.eclipse.rdf4j.sail.elasticsearch;

import java.net.InetAddress;
import java.util.concurrent.TimeUnit;

import org.elasticsearch.action.admin.cluster.health.ClusterHealthRequest;
import org.elasticsearch.action.admin.cluster.health.ClusterHealthResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.transport.TransportAddress;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.transport.client.PreBuiltTransportClient;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.Assumptions;
import org.junit.jupiter.api.BeforeAll;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.utility.DockerImageName;

@Testcontainers(disabledWithoutDocker = true)
public abstract class AbstractElasticsearchTest {

	protected static final String CLUSTER_NAME = "test";

	@Container
	public static final GenericContainer<?> elasticsearch = new GenericContainer<>(dockerImageName())
			.withEnv("discovery.type", "single-node")
			.withEnv("cluster.name", CLUSTER_NAME)
			.withEnv("ES_JAVA_OPTS",
					"-Djdk.disableLastUsageTracking=true -XX:-UseContainerSupport -Xms512m -Xmx512m")
			.withEnv("JDK_JAVA_OPTIONS",
					"-Djdk.disableLastUsageTracking=true -XX:-UseContainerSupport -Xms512m -Xmx512m")
			.withEnv("JAVA_TOOL_OPTIONS",
					"-Djdk.disableLastUsageTracking=true -XX:-UseContainerSupport -Xms512m -Xmx512m")
			.withExposedPorts(9200, 9300);

	protected static TransportClient client;

	@BeforeAll
	public static void setUpCluster() throws Exception {
		System.out.println("Setting up elasticsearch cluster");
		if (client != null) {
			return;
		}

		Assumptions.assumeTrue(elasticsearch.isRunning(),
				"Elasticsearch test container failed to start:\n" + safeLogs());

		Settings settings = Settings.builder().put("cluster.name", CLUSTER_NAME).build();

		String host = elasticsearch.getHost();
		int transportPort = elasticsearch.getMappedPort(9300);

		TransportClient transportClient = new PreBuiltTransportClient(settings)
				.addTransportAddress(new TransportAddress(InetAddress.getByName(host), transportPort));

		waitForClusterReady(transportClient);

		client = transportClient;
	}

	@AfterAll
	public static void tearDownCluster() {
		if (client != null) {
			client.close();
			client = null;
		}
	}

	private static DockerImageName dockerImageName() {
		String esVersion = System.getProperty("elasticsearch.docker.version",
				System.getProperty("elasticsearch.version", "7.15.2"));

		return DockerImageName
				.parse("docker.elastic.co/elasticsearch/elasticsearch:" + esVersion)
				.asCompatibleSubstituteFor("docker.elastic.co/elasticsearch/elasticsearch");
	}

	private static void waitForClusterReady(Client client) {
		if (!elasticsearch.isRunning()) {
			throw new IllegalStateException("Elasticsearch test container stopped before health check:\n" + safeLogs());
		}

		long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(180);
		Exception lastFailure = null;

		while (System.nanoTime() < deadline) {
			if (!elasticsearch.isRunning()) {
				throw new IllegalStateException(
						"Elasticsearch test container stopped during health check:\n" + safeLogs());
			}
			try {
				ClusterHealthRequest request = new ClusterHealthRequest()
						.waitForYellowStatus()
						.timeout(TimeValue.timeValueSeconds(1));

				ClusterHealthResponse response = client.admin().cluster().health(request).actionGet();
				if (!response.isTimedOut()) {
					return;
				}
				lastFailure = new IllegalStateException("Cluster health timed out waiting for YELLOW status");
			} catch (Exception e) {
				lastFailure = e;
			}

			try {
				Thread.sleep(100);
			} catch (InterruptedException ie) {
				Thread.currentThread().interrupt();
				throw new IllegalStateException("Interrupted while waiting for Elasticsearch test cluster", ie);
			}
		}

		throw new IllegalStateException("Timed out waiting for Elasticsearch test cluster", lastFailure);
	}

	private static String safeLogs() {
		try {
			return elasticsearch.getLogs();
		} catch (Exception e) {
			return "Unable to read container logs: " + e.getMessage();
		}
	}
}