Java tutorial
/* * Licensed under the Apache License, Version 2.0 (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.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package com.crosstreelabs.junited.elasticsearch; import com.crosstreelabs.junited.core.WrappingRule; import com.crosstreelabs.junited.elasticsearch.annotation.ElasticSetup; import com.crosstreelabs.junited.elasticsearch.annotation.ElasticSetups; import com.fasterxml.jackson.databind.ObjectMapper; import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; import java.lang.reflect.Method; import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.Map; import org.elasticsearch.action.DocumentRequest; import org.elasticsearch.action.admin.indices.delete.DeleteIndexRequest; import org.elasticsearch.action.bulk.BulkRequestBuilder; import org.elasticsearch.action.delete.DeleteRequest; import org.elasticsearch.action.index.IndexRequest; import org.elasticsearch.action.update.UpdateRequest; import org.elasticsearch.client.Client; import org.elasticsearch.common.settings.ImmutableSettings; import org.elasticsearch.common.unit.TimeValue; import org.elasticsearch.node.Node; import org.elasticsearch.node.NodeBuilder; import org.junit.runner.Description; import org.junit.runners.model.Statement; import org.slf4j.Logger; import org.slf4j.LoggerFactory; public class ElasticsearchRule extends WrappingRule { private static final Logger LOGGER = LoggerFactory.getLogger(ElasticsearchRule.class); private static final ObjectMapper MAPPER = new ObjectMapper(); private Node node; private Client client; private boolean isPersistent = true; @Override protected void before(final Statement statement, final Description description) throws Throwable { if (node == null) { ImmutableSettings.Builder settings = ImmutableSettings.settingsBuilder(); settings.put("path.data", "target/data-test"); settings.build(); NodeBuilder nb = NodeBuilder.nodeBuilder().settings(settings).local(false).client(false).data(true); node = nb.node(); } client = node.client(); List<ElasticSetup> setups = getAnnotations(description); String index = getIndex(setups, "test"); handleMapping(setups, index); handleBulkOperations(setups, index); refresh(index); } @Override protected void after(final Statement statement, final Description description) { if (!isPersistent) { client.close(); node.stop(); } } public Node getNode() { return node; } public Client getClient() { return client; } public boolean isPersistent() { return isPersistent; } public ElasticsearchRule setPersistent(final boolean isPersistent) { this.isPersistent = isPersistent; return this; } //~ Helpers ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ /** * Waits for the ES instance to return a status of yellow, which may mean * that not all shards or replicas are available, or that the index only has * a single shard or replica. */ public void waitForYellow() { client.admin().cluster().prepareHealth().setWaitForYellowStatus().setTimeout(TimeValue.timeValueSeconds(5)) .execute().actionGet(); } /** * Waits for the ES instance to return a status of green, which means that * all shards and replicas are available and that the instance is in a * stable and reliable state. */ public void waitForGreen() { client.admin().cluster().prepareHealth().setWaitForGreenStatus().setTimeout(TimeValue.timeValueSeconds(5)) .execute().actionGet(); } /** * Demand an immediate refresh of the given index, such that the index is in * an entirely consistent state given any preceding actions. * @param index The name of the index to refresh */ public void refresh(final String index) { waitForYellow(); client.admin().indices().prepareRefresh(index).execute().actionGet(); } /** * Creates a new ElasticSearch index with the given name. * @param name The name of the index to create * @param stream The mapping to create * @return The name of the newly created index. Null on failure * @throws IOException */ public String createIndex(final String name, final InputStream stream) throws IOException { waitForYellow(); if (!getClient().admin().indices().prepareCreate(name).setSource(toString(stream)).execute().actionGet() .isAcknowledged()) { throw new IOException(); } return name; } /** * Completely destroys an index with the given name. * @param name The name of the index to destroy. * @return True if destroyed, false otherwise */ public boolean destroyIndex(final String name) { waitForYellow(); return client.admin().indices().delete(new DeleteIndexRequest(name)).actionGet().isAcknowledged(); } public boolean bulkOperation(final InputStream is) throws Exception { waitForYellow(); String[] operations = toString(is).split("\n"); BulkRequestBuilder bulk = client.prepareBulk(); DocumentRequest request = null; for (String op : operations) { Map<String, Object> obj = MAPPER.readValue(op, Map.class); if (obj.containsKey("index") && obj.keySet().size() == 1) { Map<String, String> req = (Map) obj.get("index"); request = new IndexRequest(req.get("_index"), req.get("_type"), req.get("_id")); } else if (obj.containsKey("delete") && obj.keySet().size() == 1) { Map<String, String> req = (Map) obj.get("delete"); bulk.add(new DeleteRequest(req.get("_index"), req.get("_type"), req.get("_id"))); } else if (obj.containsKey("update") && obj.keySet().size() == 1) { Map<String, String> req = (Map) obj.get("update"); request = new UpdateRequest(req.get("_index"), req.get("_type"), req.get("_id")); } else if (request == null) { throw new IllegalStateException("No action to perform. Cannot add document without action"); } else if (request instanceof IndexRequest) { bulk.add(new IndexRequest((IndexRequest) request, (IndexRequest) request).source(op)); } else if (request instanceof UpdateRequest) { UpdateRequest parent = (UpdateRequest) request; bulk.add(new UpdateRequest(parent.index(), parent.type(), parent.id()).doc(op)); } } return !bulk.execute().actionGet().hasFailures(); } //~ Internal ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ protected String getIndex(final List<ElasticSetup> setups, final String fallback) { for (ElasticSetup setup : setups) { if (setup.index() != null && !setup.index().isEmpty()) { return setup.index(); } } return fallback; } protected void handleMapping(final List<ElasticSetup> setups, final String index) throws IOException { // We'll just use the first mapping found for now. We can work on // composite mapping later - Thomas for (ElasticSetup setup : setups) { for (String mapping : setup.mapping()) { InputStream is = getClass().getResourceAsStream(mapping); if (mapping == null || is == null) { throw new IllegalStateException("No mapping found"); } createIndex(index, is); return; } } } protected void handleBulkOperations(final List<ElasticSetup> setups, final String index) throws Exception { for (ElasticSetup setup : setups) { for (String file : setup.value()) { InputStream is = getClass().getResourceAsStream(file); if (is == null) { continue; } bulkOperation(is); } } } protected static List<ElasticSetup> getAnnotations(final Description description) { Class<?> testClass = description.getTestClass(); String methodName = description.getMethodName(); List<ElasticSetup> result = new ArrayList<>(); if (testClass.isAnnotationPresent(ElasticSetup.class)) { result.add(testClass.getAnnotation(ElasticSetup.class)); } if (testClass.isAnnotationPresent(ElasticSetups.class)) { result.addAll(Arrays.asList(testClass.getAnnotation(ElasticSetups.class).value())); } if (methodName == null) { return result; } try { Method method = testClass.getDeclaredMethod(methodName); if (method.isAnnotationPresent(ElasticSetup.class)) { result.add(method.getAnnotation(ElasticSetup.class)); } if (method.isAnnotationPresent(ElasticSetups.class)) { result.addAll(Arrays.asList(method.getAnnotation(ElasticSetups.class).value())); } } catch (NoSuchMethodException | SecurityException ex) { } return result; } protected static String toString(final InputStream is) throws IOException { return new String(toByteArray(is)); } protected static byte[] toByteArray(final InputStream is) throws IOException { ByteArrayOutputStream baos = new ByteArrayOutputStream(); int length = 0; byte[] buf = new byte[256]; while ((length = is.read(buf)) > 0) { baos.write(buf, 0, length); } return baos.toByteArray(); } }