Skip to content

Commit 17a3568

Browse files
author
Lingxi Chen
committed
trying to fix Unknown feature type
Signed-off-by: Lingxi Chen <lingxich@amazon.com>
1 parent fceb87a commit 17a3568

File tree

1 file changed

+41
-105
lines changed

1 file changed

+41
-105
lines changed

‎plugins/workload-management/src/internalClusterTest/java/org/opensearch/plugin/wlm/WlmAutoTaggingIT.java‎

Lines changed: 41 additions & 105 deletions
Original file line numberDiff line numberDiff line change
@@ -10,22 +10,21 @@
1010

1111
import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
1212

13+
import org.junit.After;
14+
import org.junit.Before;
15+
import org.opensearch.action.ActionRequest;
16+
import org.opensearch.action.ActionType;
1317
import org.opensearch.action.support.PlainActionFuture;
1418
import org.opensearch.action.support.clustermanager.AcknowledgedResponse;
19+
import org.opensearch.common.action.ActionFuture;
1520
import org.opensearch.common.settings.Settings;
1621
import org.opensearch.core.action.ActionListener;
22+
import org.opensearch.core.action.ActionResponse;
1723
import org.opensearch.plugin.wlm.rule.WorkloadGroupFeatureType;
24+
import org.opensearch.plugins.ActionPlugin;
1825
import org.opensearch.plugins.Plugin;
19-
import org.opensearch.rule.RuleAttribute;
20-
import org.opensearch.rule.RulePersistenceService;
21-
import org.opensearch.rule.RuleRoutingService;
22-
import org.opensearch.rule.action.CreateRuleRequest;
23-
import org.opensearch.rule.action.CreateRuleResponse;
24-
import org.opensearch.rule.action.DeleteRuleRequest;
25-
import org.opensearch.rule.action.GetRuleRequest;
26-
import org.opensearch.rule.action.GetRuleResponse;
27-
import org.opensearch.rule.action.UpdateRuleRequest;
28-
import org.opensearch.rule.action.UpdateRuleResponse;
26+
import org.opensearch.rule.*;
27+
import org.opensearch.rule.action.*;
2928
import org.opensearch.rule.autotagging.Attribute;
3029
import org.opensearch.rule.autotagging.AutoTaggingRegistry;
3130
import org.opensearch.rule.autotagging.FeatureType;
@@ -61,6 +60,8 @@ public static Collection<Object[]> parameters() {
6160
@Override
6261
protected Collection<Class<? extends Plugin>> nodePlugins() {
6362
List<Class<? extends Plugin>> plugins = new ArrayList<>(super.nodePlugins());
63+
plugins.add(TestRuleFrameworkPlugin.class);
64+
plugins.add(WorkloadManagementPlugin.class);
6465
return plugins;
6566
}
6667

@@ -69,116 +70,41 @@ public static void registerWorkloadGroupFeatureType() {
6970
AutoTaggingRegistry.registerFeatureType(new WorkloadGroupFeatureType(featureValue -> {}));
7071
}
7172

72-
public void testCreateRuleWithSingleAttribute() throws Exception {
73-
FeatureType featureType = AutoTaggingRegistry.getFeatureType("workload_group");
74-
Map<Attribute, Set<String>> attributeMap = Map.of(RuleAttribute.INDEX_PATTERN, Set.of("logs-*"));
75-
76-
Rule rule = new Rule("mock_create", "desc", attributeMap, featureType, "feature1", "2025-06-20T00:00:00.000Z");
77-
CreateRuleRequest request = new CreateRuleRequest(rule);
78-
PlainActionFuture<CreateRuleResponse> future = PlainActionFuture.newFuture();
73+
private static final String RULE_ID = "test_rule_id";
7974

80-
new TestRulePersistenceService().createRule(request, future);
81-
CreateRuleResponse response = future.get();
8275

83-
assertEquals("mock_create", response.getRule().getId());
84-
assertEquals("feature1", response.getRule().getFeatureValue());
85-
}
86-
87-
public void testGetExistingRule() throws Exception {
88-
FeatureType featureType = AutoTaggingRegistry.getFeatureType("workload_group");
89-
Map<Attribute, Set<String>> attributeMap = Map.of(RuleAttribute.INDEX_PATTERN, Set.of("logs-*"));
90-
91-
GetRuleRequest request = new GetRuleRequest("mock_id", attributeMap, "token", featureType);
92-
PlainActionFuture<GetRuleResponse> future = PlainActionFuture.newFuture();
93-
94-
new TestRulePersistenceService().getRule(request, future);
95-
GetRuleResponse response = future.get();
96-
97-
assertEquals(1, response.getRules().size());
98-
assertEquals("mock_id", response.getRules().get(0).getId());
99-
}
10076

101-
public void testUpdateRuleValueChange() throws Exception {
77+
public void testCreateAutoTaggingRule() throws Exception {
10278
FeatureType featureType = AutoTaggingRegistry.getFeatureType("workload_group");
103-
Map<Attribute, Set<String>> attributeMap = Map.of(RuleAttribute.INDEX_PATTERN, Set.of("logs-*"));
104-
105-
Rule updatedRule = new Rule("mock_id", "updated desc", attributeMap, featureType, "new_feature_value", "2025-06-20T00:00:00.000Z");
106-
UpdateRuleRequest request = new UpdateRuleRequest(
107-
updatedRule.getId(),
108-
updatedRule.getDescription(),
109-
updatedRule.getAttributeMap(),
110-
updatedRule.getFeatureValue(),
111-
updatedRule.getFeatureType()
79+
Map<Attribute, Set<String>> attributes = Map.of(
80+
RuleAttribute.INDEX_PATTERN, Set.of("logs-*")
11281
);
11382

114-
PlainActionFuture<UpdateRuleResponse> future = PlainActionFuture.newFuture();
115-
new TestRulePersistenceService().updateRule(request, future);
116-
UpdateRuleResponse response = future.get();
117-
118-
assertEquals("new_feature_value", response.getRule().getFeatureValue());
119-
}
120-
121-
public void testDeleteRule() throws Exception {
122-
FeatureType featureType = AutoTaggingRegistry.getFeatureType("workload_group");
123-
DeleteRuleRequest request = new DeleteRuleRequest("mock_id", featureType);
124-
PlainActionFuture<AcknowledgedResponse> future = PlainActionFuture.newFuture();
125-
126-
new TestRulePersistenceService().deleteRule(request, future);
127-
AcknowledgedResponse response = future.get();
128-
129-
assertTrue(response.isAcknowledged());
130-
}
131-
132-
public void testCreateRuleWithMultiplePatterns() throws Exception {
133-
FeatureType featureType = AutoTaggingRegistry.getFeatureType("workload_group");
134-
Map<Attribute, Set<String>> attributeMap = Map.of(RuleAttribute.INDEX_PATTERN, Set.of("logs-*", "metrics-*", "events-*"));
135-
136-
Rule rule = new Rule("multi_pattern", "desc", attributeMap, featureType, "value", "2025-06-20T00:00:00.000Z");
137-
CreateRuleRequest request = new CreateRuleRequest(rule);
138-
PlainActionFuture<CreateRuleResponse> future = PlainActionFuture.newFuture();
139-
140-
new TestRulePersistenceService().createRule(request, future);
141-
CreateRuleResponse response = future.get();
142-
143-
assertEquals("multi_pattern", response.getRule().getId());
144-
assertTrue(response.getRule().getAttributeMap().get(RuleAttribute.INDEX_PATTERN).contains("events-*"));
145-
}
146-
147-
public void testCreateRuleWithExistingId() throws Exception {
148-
FeatureType featureType = AutoTaggingRegistry.getFeatureType("workload_group");
149-
Rule rule1 = new Rule(
150-
"duplicate_id",
151-
"desc1",
152-
Map.of(RuleAttribute.INDEX_PATTERN, Set.of("logs-*")),
83+
Rule rule = new Rule(
84+
RULE_ID,
85+
"Integration test rule",
86+
attributes,
15387
featureType,
154-
"v1",
88+
"wlm_group_id_123",
15589
"2025-06-20T00:00:00.000Z"
15690
);
157-
Rule rule2 = new Rule(
158-
"duplicate_id",
159-
"desc2",
160-
Map.of(RuleAttribute.INDEX_PATTERN, Set.of("metrics-*")),
161-
featureType,
162-
"v2",
163-
"2025-06-20T00:00:00.000Z"
164-
);
165-
166-
CreateRuleRequest request1 = new CreateRuleRequest(rule1);
167-
CreateRuleRequest request2 = new CreateRuleRequest(rule2);
168-
PlainActionFuture<CreateRuleResponse> future1 = PlainActionFuture.newFuture();
169-
PlainActionFuture<CreateRuleResponse> future2 = PlainActionFuture.newFuture();
17091

171-
new TestRulePersistenceService().createRule(request1, future1);
172-
new TestRulePersistenceService().createRule(request2, future2);
92+
CreateRuleRequest createRequest = new CreateRuleRequest(rule);
93+
CreateRuleResponse response = client().execute(CreateRuleAction.INSTANCE, createRequest).get();
17394

174-
assertEquals("duplicate_id", future1.get().getRule().getId());
175-
assertEquals("duplicate_id", future2.get().getRule().getId());
95+
assertNotNull(response);
96+
assertEquals(RULE_ID, response.getRule().getId());
97+
assertEquals("wlm_group_id_123", response.getRule().getFeatureValue());
17698
}
17799

178-
public static class TestRuleFrameworkPlugin extends Plugin implements RuleFrameworkExtension {
100+
public static class TestRuleFrameworkPlugin extends Plugin implements RuleFrameworkExtension, ActionPlugin {
101+
public TestRuleFrameworkPlugin() {
102+
super();
103+
}
104+
179105
@Override
180106
public Supplier<FeatureType> getFeatureTypeSupplier() {
181-
return () -> new WorkloadGroupFeatureType(featureValue -> {});
107+
return () -> new WorkloadGroupFeatureType(v -> {});
182108
}
183109

184110
@Override
@@ -190,6 +116,16 @@ public Supplier<RulePersistenceService> getRulePersistenceServiceSupplier() {
190116
public Supplier<RuleRoutingService> getRuleRoutingServiceSupplier() {
191117
return TestRuleRoutingService::new;
192118
}
119+
120+
@Override
121+
public List<ActionPlugin.ActionHandler<? extends ActionRequest, ? extends ActionResponse>> getActions() {
122+
return List.of(
123+
new ActionPlugin.ActionHandler<>(CreateRuleAction.INSTANCE, TransportCreateRuleAction.class),
124+
new ActionPlugin.ActionHandler<>(GetRuleAction.INSTANCE, TransportGetRuleAction.class),
125+
new ActionPlugin.ActionHandler<>(UpdateRuleAction.INSTANCE, TransportUpdateRuleAction.class),
126+
new ActionPlugin.ActionHandler<>(DeleteRuleAction.INSTANCE, TransportDeleteRuleAction.class)
127+
);
128+
}
193129
}
194130

195131
public static class TestRulePersistenceService implements RulePersistenceService {

0 commit comments

Comments
 (0)