Skip to content

Commit 3ad0855

Browse files
author
John Plaisted
authored
feat: finish off implementation ElasticsearchIntegrationTestExtension (#41)
1 parent 57b8103 commit 3ad0855

File tree

3 files changed

+219
-11
lines changed

3 files changed

+219
-11
lines changed

‎testing/elasticsearch-dao-integ-testing/src/main/java/com/linkedin/metadata/testing/ElasticsearchIntegrationTestExtension.java

Lines changed: 165 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,26 @@
11
package com.linkedin.metadata.testing;
22

3+
import com.linkedin.metadata.testing.annotations.SearchIndexMappings;
4+
import com.linkedin.metadata.testing.annotations.SearchIndexSettings;
5+
import com.linkedin.metadata.testing.annotations.SearchIndexType;
6+
import java.io.IOException;
7+
import java.lang.reflect.Constructor;
8+
import java.lang.reflect.Field;
9+
import java.net.URL;
10+
import java.util.ArrayList;
11+
import java.util.List;
12+
import java.util.function.Function;
13+
import java.util.stream.Collectors;
14+
import javax.annotation.Nonnull;
15+
import org.apache.commons.io.IOUtils;
316
import org.junit.jupiter.api.extension.AfterAllCallback;
417
import org.junit.jupiter.api.extension.AfterEachCallback;
518
import org.junit.jupiter.api.extension.BeforeAllCallback;
619
import org.junit.jupiter.api.extension.BeforeEachCallback;
20+
import org.junit.jupiter.api.extension.ExtensionConfigurationException;
721
import org.junit.jupiter.api.extension.ExtensionContext;
22+
import org.junit.platform.commons.util.ClassFilter;
23+
import org.junit.platform.commons.util.ReflectionUtils;
824

925

1026
/**
@@ -14,24 +30,165 @@
1430
*/
1531
final class ElasticsearchIntegrationTestExtension
1632
implements BeforeAllCallback, AfterAllCallback, BeforeEachCallback, AfterEachCallback {
33+
private static final ExtensionContext.Namespace NAMESPACE =
34+
ExtensionContext.Namespace.create(ElasticsearchIntegrationTestExtension.class);
35+
36+
private final static String CONTAINER_FACTORY = "containerFactory";
37+
private final static String CONNECTION = "connection";
38+
private final static String STATIC_INDICIES = "staticIndicies";
39+
private final static String INDICIES = "indicies";
1740

1841
@Override
19-
public void afterAll(ExtensionContext context) throws Exception {
20-
// TODO
42+
public void beforeAll(ExtensionContext context) throws Exception {
43+
final Class<?> testClass = context.getTestClass()
44+
.orElseThrow(() -> new ExtensionConfigurationException(
45+
"ElasticSearchIntegrationTestExtension is only supported for classes."));
46+
47+
final ExtensionContext.Store store = context.getStore(NAMESPACE);
48+
49+
final ElasticsearchContainerFactory factory = getContainerFactory();
50+
final ElasticsearchConnection connection = factory.start();
51+
52+
store.put(CONTAINER_FACTORY, factory);
53+
store.put(CONNECTION, connection);
54+
55+
final List<Field> fields = ReflectionUtils.findFields(testClass, field -> {
56+
return ReflectionUtils.isStatic(field) && ReflectionUtils.isPublic(field) && ReflectionUtils.isNotFinal(field)
57+
&& SearchIndex.class.isAssignableFrom(field.getType());
58+
}, ReflectionUtils.HierarchyTraversalMode.TOP_DOWN);
59+
60+
final SearchIndexFactory indexFactory = new SearchIndexFactory(connection);
61+
final List<SearchIndex<?>> indices = createIndices(indexFactory, context.getRequiredTestClass(), fields,
62+
fieldName -> String.format("%s_%s_%s", fieldName, testClass.getSimpleName(), System.currentTimeMillis()));
63+
store.put(STATIC_INDICIES, indices);
64+
}
65+
66+
private List<SearchIndex<?>> createIndices(@Nonnull SearchIndexFactory indexFactory, @Nonnull Object testInstance,
67+
@Nonnull List<Field> fields, @Nonnull Function<String, String> nameFn) throws Exception {
68+
final List<SearchIndex<?>> indices = new ArrayList<>();
69+
70+
for (Field field : fields) {
71+
final SearchIndexType searchIndexType = field.getAnnotation(SearchIndexType.class);
72+
73+
if (searchIndexType == null) {
74+
throw new IllegalStateException(
75+
String.format("Field `%s` must be annotated with `SearchIndexType`.", field.getName()));
76+
}
77+
78+
final String indexName = nameFn.apply(field.getName()).replaceAll("^_*", "").toLowerCase();
79+
80+
final SearchIndexSettings settings = field.getAnnotation(SearchIndexSettings.class);
81+
final String settingsJson = settings == null ? null : loadResource(testInstance.getClass(), settings.value());
82+
83+
final SearchIndexMappings mappings = field.getAnnotation(SearchIndexMappings.class);
84+
final String mappingsJson = mappings == null ? null : loadResource(testInstance.getClass(), mappings.value());
85+
86+
final SearchIndex<?> index =
87+
indexFactory.createIndex(searchIndexType.value(), indexName, settingsJson, mappingsJson);
88+
field.set(testInstance, index);
89+
indices.add(index);
90+
}
91+
92+
return indices;
93+
}
94+
95+
private String loadResource(@Nonnull Class<?> testClass, @Nonnull String resource) throws IOException {
96+
final URL resourceUrl = testClass.getResource(resource);
97+
if (resourceUrl == null) {
98+
throw new IllegalArgumentException(String.format("Resource `%s` not found.", resource));
99+
}
100+
return IOUtils.toString(resourceUrl);
101+
}
102+
103+
@Nonnull
104+
private Class<?> findContainerFactoryClass() {
105+
final List<Class<?>> classes = ReflectionUtils.findAllClassesInPackage("com.linkedin.metadata.testing",
106+
ClassFilter.of(clazz -> clazz.isAnnotationPresent(ElasticsearchContainerFactory.Implementation.class)));
107+
108+
if (classes.size() == 0) {
109+
throw new IllegalStateException("Could not find any ElasticsearchContainerFactory implementations.");
110+
}
111+
112+
if (classes.size() > 1) {
113+
throw new IllegalStateException(
114+
String.format("Found %s ElasticsearchContainerFactory implementations, expected 1. Found %s.", classes.size(),
115+
String.join(", ", classes.stream().map(Class::getSimpleName).collect(Collectors.toList()))));
116+
}
117+
118+
return classes.get(0);
119+
}
120+
121+
@Nonnull
122+
private ElasticsearchContainerFactory getContainerFactory() throws Exception {
123+
final Class<?> clazz = findContainerFactoryClass();
124+
125+
if (!ElasticsearchContainerFactory.class.isAssignableFrom(clazz)) {
126+
throw new IllegalStateException(String.format(
127+
"Provided class `%s` to ElasticsearchIntegrationTest, but did not inherit from "
128+
+ "ElasticsearchContainerFactory.", clazz.toString()));
129+
}
130+
131+
Constructor<?> constructor;
132+
try {
133+
constructor = clazz.getConstructor();
134+
} catch (NoSuchMethodException e) {
135+
throw new NoSuchMethodException(String.format(
136+
"Expected ElasticsearchContainerFactory, `%s`, to have a default, public, constructor but found none.",
137+
clazz.toString()));
138+
}
139+
140+
return (ElasticsearchContainerFactory) constructor.newInstance();
21141
}
22142

23143
@Override
24-
public void afterEach(ExtensionContext context) throws Exception {
25-
// TODO
144+
public void beforeEach(ExtensionContext context) throws Exception {
145+
final ExtensionContext.Store store = context.getStore(NAMESPACE);
146+
final ElasticsearchConnection connection = store.get(CONNECTION, ElasticsearchConnection.class);
147+
148+
final List<Field> fields = ReflectionUtils.findFields(context.getRequiredTestClass(), field -> {
149+
return ReflectionUtils.isNotStatic(field) && ReflectionUtils.isPublic(field) && ReflectionUtils.isNotFinal(field)
150+
&& SearchIndex.class.isAssignableFrom(field.getType());
151+
}, ReflectionUtils.HierarchyTraversalMode.TOP_DOWN);
152+
153+
final SearchIndexFactory indexFactory = new SearchIndexFactory(connection);
154+
final List<SearchIndex<?>> indices = createIndices(indexFactory, context.getRequiredTestInstance(), fields,
155+
fieldName -> String.format("%s_%s_%s_%s", fieldName, context.getRequiredTestMethod().getName(),
156+
context.getRequiredTestClass().getSimpleName(), System.currentTimeMillis()));
157+
store.put(INDICIES, indices);
26158
}
27159

160+
@SuppressWarnings("unchecked")
28161
@Override
29-
public void beforeAll(ExtensionContext context) throws Exception {
30-
// TODO
162+
public void afterAll(ExtensionContext context) throws Exception {
163+
final ExtensionContext.Store store = context.getStore(NAMESPACE);
164+
165+
final List<SearchIndex<?>> indices = (List<SearchIndex<?>>) store.get(STATIC_INDICIES, List.class);
166+
final ElasticsearchConnection connection = store.get(CONNECTION, ElasticsearchConnection.class);
167+
168+
cleanUp(connection, indices);
169+
170+
// don't need to close the factory since it implements CloseableResource, junit will close it since it is in the
171+
// store
31172
}
32173

174+
@SuppressWarnings("unchecked")
33175
@Override
34-
public void beforeEach(ExtensionContext context) throws Exception {
35-
// TODO
176+
public void afterEach(ExtensionContext context) throws Exception {
177+
final ExtensionContext.Store store = context.getStore(NAMESPACE);
178+
179+
final List<SearchIndex<?>> indices = (List<SearchIndex<?>>) store.get(INDICIES, List.class);
180+
final ElasticsearchConnection connection = store.get(CONNECTION, ElasticsearchConnection.class);
181+
182+
if (indices != null) {
183+
cleanUp(connection, indices);
184+
}
185+
}
186+
187+
private void cleanUp(@Nonnull ElasticsearchConnection connection, @Nonnull List<SearchIndex<?>> indices) {
188+
for (SearchIndex<?> i : indices) {
189+
connection.getTransportClient().admin().indices().prepareDelete(i.getName()).get();
190+
}
191+
192+
indices.clear();
36193
}
37194
}

‎testing/elasticsearch-dao-integ-testing/src/main/java/com/linkedin/metadata/testing/SearchIndexFactory.java

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
package com.linkedin.metadata.testing;
2+
3+
import com.linkedin.data.template.RecordTemplate;
4+
import java.util.concurrent.ExecutionException;
5+
import javax.annotation.Nonnull;
6+
import javax.annotation.Nullable;
7+
import org.elasticsearch.action.admin.indices.create.CreateIndexRequest;
8+
import org.elasticsearch.common.xcontent.XContentType;
9+
10+
11+
/**
12+
* Factory to create {@link SearchIndex} instances for testing.
13+
*/
14+
final class SearchIndexFactory {
15+
private final ElasticsearchConnection _connection;
16+
17+
SearchIndexFactory(@Nonnull ElasticsearchConnection connection) {
18+
_connection = connection;
19+
}
20+
21+
/**
22+
* Creates a search index to read / write the given document type for testing.
23+
*
24+
* <p>This will create an index on the Elasticsearch instance with a unique name.
25+
*
26+
* @param documentClass the document type
27+
* @param name the name to use for the index
28+
*/
29+
public <DOCUMENT extends RecordTemplate> SearchIndex<DOCUMENT> createIndex(@Nonnull Class<DOCUMENT> documentClass,
30+
@Nonnull String name, @Nullable String settingsJson, @Nullable String mappingsJson) {
31+
final CreateIndexRequest createIndexRequest = new CreateIndexRequest(name);
32+
33+
if (settingsJson != null) {
34+
createIndexRequest.settings(settingsJson, XContentType.JSON);
35+
}
36+
37+
if (mappingsJson != null) {
38+
// TODO
39+
createIndexRequest.mapping("doc", mappingsJson, XContentType.JSON);
40+
}
41+
42+
try {
43+
_connection.getTransportClient().admin().indices().create(createIndexRequest).get();
44+
} catch (InterruptedException | ExecutionException e) {
45+
throw new RuntimeException(e);
46+
}
47+
48+
return new SearchIndex<>(documentClass, _connection, name);
49+
}
50+
}

‎testing/elasticsearch-dao-integ-testing/src/main/java/com/linkedin/metadata/testing/annotations/SearchIndexType.java

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.linkedin.metadata.testing.annotations;
22

33
import com.linkedin.data.template.RecordTemplate;
4+
import com.linkedin.metadata.testing.SearchIndex;
45
import java.lang.annotation.ElementType;
56
import java.lang.annotation.Retention;
67
import java.lang.annotation.RetentionPolicy;
@@ -9,17 +10,17 @@
910

1011

1112
/**
12-
* Annotates the given {@link com.linkedin.metadata.testing.SearchIndex} field with the document type.
13+
* Annotates the given {@link SearchIndex} field with the document type.
1314
*
14-
* <p>Required annotation for {@link com.linkedin.metadata.testing.SearchIndex} instances in tests.</p>
15+
* <p>Required annotation for {@link SearchIndex} instances in tests.</p>
1516
*/
1617
@Target(ElementType.FIELD)
1718
@Retention(RetentionPolicy.RUNTIME)
1819
public @interface SearchIndexType {
1920
/**
2021
* The search document class for this index.
2122
*
22-
* <p>Used to create an instance of the {@link com.linkedin.metadata.testing.SearchIndex} during testing.
23+
* <p>Used to create an instance of the {@link SearchIndex} during testing.
2324
*/
2425
@Nonnull
2526
Class<? extends RecordTemplate> value();

0 commit comments

Comments
 (0)