Skip to content

Commit a8c21a4

Browse files
[ML] Add integration test for model plots (#30359)
Relates #30004
1 parent a57512c commit a8c21a4

File tree

2 files changed

+164
-4
lines changed

2 files changed

+164
-4
lines changed

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/job/config/ModelPlotConfig.java

-4
Original file line numberDiff line numberDiff line change
@@ -53,10 +53,6 @@ public ModelPlotConfig() {
5353
this(true, null);
5454
}
5555

56-
public ModelPlotConfig(boolean enabled) {
57-
this(false, null);
58-
}
59-
6056
public ModelPlotConfig(boolean enabled, String terms) {
6157
this.enabled = enabled;
6258
this.terms = terms;
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License;
4+
* you may not use this file except in compliance with the Elastic License.
5+
*/
6+
package org.elasticsearch.xpack.ml.integration;
7+
8+
import org.elasticsearch.action.bulk.BulkRequestBuilder;
9+
import org.elasticsearch.action.bulk.BulkResponse;
10+
import org.elasticsearch.action.index.IndexRequest;
11+
import org.elasticsearch.action.search.SearchResponse;
12+
import org.elasticsearch.action.support.WriteRequest;
13+
import org.elasticsearch.common.unit.TimeValue;
14+
import org.elasticsearch.index.query.QueryBuilders;
15+
import org.elasticsearch.search.aggregations.AggregationBuilders;
16+
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
17+
import org.elasticsearch.xpack.core.ml.datafeed.DatafeedConfig;
18+
import org.elasticsearch.xpack.core.ml.job.config.AnalysisConfig;
19+
import org.elasticsearch.xpack.core.ml.job.config.DataDescription;
20+
import org.elasticsearch.xpack.core.ml.job.config.Detector;
21+
import org.elasticsearch.xpack.core.ml.job.config.Job;
22+
import org.elasticsearch.xpack.core.ml.job.config.ModelPlotConfig;
23+
import org.junit.After;
24+
import org.junit.Before;
25+
26+
import java.util.Arrays;
27+
import java.util.List;
28+
import java.util.Set;
29+
import java.util.stream.Collectors;
30+
31+
import static org.hamcrest.Matchers.containsInAnyOrder;
32+
import static org.hamcrest.Matchers.equalTo;
33+
import static org.hamcrest.Matchers.is;
34+
35+
public class ModelPlotsIT extends MlNativeAutodetectIntegTestCase {
36+
37+
private static final String DATA_INDEX = "model-plots-test-data";
38+
private static final String DATA_TYPE = "doc";
39+
40+
@Before
41+
public void setUpData() {
42+
client().admin().indices().prepareCreate(DATA_INDEX)
43+
.addMapping(DATA_TYPE, "time", "type=date,format=epoch_millis", "user", "type=keyword")
44+
.get();
45+
46+
List<String> users = Arrays.asList("user_1", "user_2", "user_3");
47+
48+
// We are going to create data for last day
49+
long nowMillis = System.currentTimeMillis();
50+
int totalBuckets = 24;
51+
BulkRequestBuilder bulkRequestBuilder = client().prepareBulk();
52+
for (int bucket = 0; bucket < totalBuckets; bucket++) {
53+
long timestamp = nowMillis - TimeValue.timeValueHours(totalBuckets - bucket).getMillis();
54+
for (String user : users) {
55+
IndexRequest indexRequest = new IndexRequest(DATA_INDEX, DATA_TYPE);
56+
indexRequest.source("time", timestamp, "user", user);
57+
bulkRequestBuilder.add(indexRequest);
58+
}
59+
}
60+
61+
BulkResponse bulkResponse = bulkRequestBuilder
62+
.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
63+
.get();
64+
assertThat(bulkResponse.hasFailures(), is(false));
65+
}
66+
67+
@After
68+
public void tearDownData() {
69+
client().admin().indices().prepareDelete(DATA_INDEX).get();
70+
cleanUp();
71+
}
72+
73+
public void testPartitionFieldWithoutTerms() throws Exception {
74+
Job.Builder job = jobWithPartitionUser("model-plots-it-test-partition-field-without-terms");
75+
job.setModelPlotConfig(new ModelPlotConfig());
76+
putJob(job);
77+
String datafeedId = job.getId() + "-feed";
78+
DatafeedConfig datafeed = newDatafeed(datafeedId, job.getId());
79+
registerDatafeed(datafeed);
80+
putDatafeed(datafeed);
81+
openJob(job.getId());
82+
startDatafeed(datafeedId, 0, System.currentTimeMillis());
83+
waitUntilJobIsClosed(job.getId());
84+
85+
assertThat(getBuckets(job.getId()).size(), equalTo(23));
86+
Set<String> modelPlotTerms = modelPlotTerms(job.getId(), "partition_field_value");
87+
assertThat(modelPlotTerms, containsInAnyOrder("user_1", "user_2", "user_3"));
88+
}
89+
90+
public void testPartitionFieldWithTerms() throws Exception {
91+
Job.Builder job = jobWithPartitionUser("model-plots-it-test-partition-field-with-terms");
92+
job.setModelPlotConfig(new ModelPlotConfig(true, "user_2,user_3"));
93+
putJob(job);
94+
String datafeedId = job.getId() + "-feed";
95+
DatafeedConfig datafeed = newDatafeed(datafeedId, job.getId());
96+
registerDatafeed(datafeed);
97+
putDatafeed(datafeed);
98+
openJob(job.getId());
99+
startDatafeed(datafeedId, 0, System.currentTimeMillis());
100+
waitUntilJobIsClosed(job.getId());
101+
102+
assertThat(getBuckets(job.getId()).size(), equalTo(23));
103+
Set<String> modelPlotTerms = modelPlotTerms(job.getId(), "partition_field_value");
104+
assertThat(modelPlotTerms, containsInAnyOrder("user_2", "user_3"));
105+
}
106+
107+
public void testByFieldWithTerms() throws Exception {
108+
Job.Builder job = jobWithByUser("model-plots-it-test-by-field-with-terms");
109+
job.setModelPlotConfig(new ModelPlotConfig(true, "user_2,user_3"));
110+
putJob(job);
111+
String datafeedId = job.getId() + "-feed";
112+
DatafeedConfig datafeed = newDatafeed(datafeedId, job.getId());
113+
registerDatafeed(datafeed);
114+
putDatafeed(datafeed);
115+
openJob(job.getId());
116+
startDatafeed(datafeedId, 0, System.currentTimeMillis());
117+
waitUntilJobIsClosed(job.getId());
118+
119+
assertThat(getBuckets(job.getId()).size(), equalTo(23));
120+
Set<String> modelPlotTerms = modelPlotTerms(job.getId(), "by_field_value");
121+
assertThat(modelPlotTerms, containsInAnyOrder("user_2", "user_3"));
122+
}
123+
124+
private static Job.Builder jobWithPartitionUser(String id) {
125+
Detector.Builder detector = new Detector.Builder();
126+
detector.setFunction("count");
127+
detector.setPartitionFieldName("user");
128+
return newJobBuilder(id, detector.build());
129+
}
130+
131+
private static Job.Builder jobWithByUser(String id) {
132+
Detector.Builder detector = new Detector.Builder();
133+
detector.setFunction("count");
134+
detector.setByFieldName("user");
135+
return newJobBuilder(id, detector.build());
136+
}
137+
138+
private static Job.Builder newJobBuilder(String id, Detector detector) {
139+
AnalysisConfig.Builder analysisConfig = new AnalysisConfig.Builder(Arrays.asList(detector));
140+
analysisConfig.setBucketSpan(TimeValue.timeValueHours(1));
141+
DataDescription.Builder dataDescription = new DataDescription.Builder();
142+
dataDescription.setTimeField("time");
143+
Job.Builder jobBuilder = new Job.Builder(id);
144+
jobBuilder.setAnalysisConfig(analysisConfig);
145+
jobBuilder.setDataDescription(dataDescription);
146+
return jobBuilder;
147+
}
148+
149+
private static DatafeedConfig newDatafeed(String datafeedId, String jobId) {
150+
DatafeedConfig.Builder datafeedConfig = new DatafeedConfig.Builder(datafeedId, jobId);
151+
datafeedConfig.setIndices(Arrays.asList(DATA_INDEX));
152+
return datafeedConfig.build();
153+
}
154+
155+
private Set<String> modelPlotTerms(String jobId, String fieldName) {
156+
SearchResponse searchResponse = client().prepareSearch(".ml-anomalies-" + jobId)
157+
.setQuery(QueryBuilders.termQuery("result_type", "model_plot"))
158+
.addAggregation(AggregationBuilders.terms("model_plot_terms").field(fieldName))
159+
.get();
160+
161+
Terms aggregation = searchResponse.getAggregations().get("model_plot_terms");
162+
return aggregation.getBuckets().stream().map(agg -> agg.getKeyAsString()).collect(Collectors.toSet());
163+
}
164+
}

0 commit comments

Comments
 (0)