Skip to content

Commit 224715f

Browse files
committed
feat: custom mini testing module selection for cyborg
1 parent 10c0a05 commit 224715f

File tree

11 files changed

+236
-28
lines changed

11 files changed

+236
-28
lines changed

apps/dashboard/src/main/java/com/akto/action/testing/StartTestAction.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,7 @@ private TestingRun createTestingRun(int scheduleTimestamp, int periodInSeconds)
159159

160160
return new TestingRun(scheduleTimestamp, user.getLogin(),
161161
testingEndpoints, testIdConfig, State.SCHEDULED, periodInSeconds, testName, this.testRunTime,
162-
this.maxConcurrentRequests, this.sendSlackAlert);
162+
this.maxConcurrentRequests, this.sendSlackAlert, null);
163163
}
164164

165165
private List<String> selectedTests;

apps/database-abstractor/src/main/java/com/akto/action/DbAction.java

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,7 @@ public void setIssuesIds(BasicDBList issuesIds) {
204204
Bson completedUpdate;
205205
int totalApiCount;
206206
boolean hybridTestingEnabled;
207+
String miniTestingServiceName;
207208
TestingRun testingRun;
208209
TestingRunConfig testingRunConfig;
209210
Boolean exists;
@@ -1303,7 +1304,7 @@ public String createTRRSummaryIfAbsent() {
13031304

13041305
public String findPendingTestingRun() {
13051306
try {
1306-
testingRun = DbLayer.findPendingTestingRun(delta);
1307+
testingRun = DbLayer.findPendingTestingRun(delta, miniTestingServiceName);
13071308
if (testingRun != null) {
13081309
/*
13091310
* There is a db call involved for collectionWiseTestingEndpoints, thus this hack.
@@ -1323,7 +1324,7 @@ public String findPendingTestingRun() {
13231324

13241325
public String findPendingTestingRunResultSummary() {
13251326
try {
1326-
trrs = DbLayer.findPendingTestingRunResultSummary(now, delta);
1327+
trrs = DbLayer.findPendingTestingRunResultSummary(now, delta, miniTestingServiceName);
13271328
if (trrs != null) {
13281329
trrs.setTestingRunHexId(trrs.getTestingRunId().toHexString());
13291330
}
@@ -1935,6 +1936,15 @@ public String modifyHybridTestingSetting() {
19351936
return Action.SUCCESS.toUpperCase();
19361937
}
19371938

1939+
public String modifyHybridTestingSettingWithCustomName() {
1940+
try {
1941+
DbLayer.modifyHybridTestingSettingWithCustomName(hybridTestingEnabled, miniTestingServiceName);
1942+
} catch (Exception e) {
1943+
return Action.ERROR.toUpperCase();
1944+
}
1945+
return Action.SUCCESS.toUpperCase();
1946+
}
1947+
19381948
public String insertTestingLog() {
19391949
try {
19401950
int accId = Context.accountId.get();
@@ -3356,4 +3366,11 @@ public void setOperator(String operator) {
33563366
this.operator = operator;
33573367
}
33583368

3369+
public String getMiniTestingServiceName() {
3370+
return miniTestingServiceName;
3371+
}
3372+
3373+
public void setMiniTestingServiceName(String miniTestingServiceName) {
3374+
this.miniTestingServiceName = miniTestingServiceName;
3375+
}
33593376
}

apps/database-abstractor/src/main/resources/struts.xml

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1266,6 +1266,17 @@
12661266
</result>
12671267
</action>
12681268

1269+
<action name="api/modifyTestingSettingWithCustomName" class="com.akto.action.DbAction" method="modifyHybridTestingSettingWithCustomName">
1270+
<interceptor-ref name="json"/>
1271+
<interceptor-ref name="defaultStack" />
1272+
<result name="SUCCESS" type="json"/>
1273+
<result name="ERROR" type="json">
1274+
<param name="statusCode">422</param>
1275+
<param name="ignoreHierarchy">false</param>
1276+
<param name="includeProperties">^actionErrors.*</param>
1277+
</result>
1278+
</action>
1279+
12691280
<action name="api/updateIssueCountAndStateInSummary" class="com.akto.action.DbAction" method="updateIssueCountAndStateInSummary">
12701281
<interceptor-ref name="json"/>
12711282
<interceptor-ref name="defaultStack" />

apps/mini-testing/src/main/java/com/akto/testing/Main.java

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package com.akto.testing;
22

3+
import com.akto.RuntimeMode;
34
import com.akto.billing.UsageMetricUtils;
45
import com.akto.dao.context.Context;
56
import com.akto.data_actor.DataActor;
@@ -55,6 +56,8 @@ public class Main {
5556
public static boolean SKIP_SSRF_CHECK = ("true".equalsIgnoreCase(System.getenv("SKIP_SSRF_CHECK")) || !DashboardMode.isSaasDeployment());
5657
public static final boolean IS_SAAS = "true".equalsIgnoreCase(System.getenv("IS_SAAS"));
5758

59+
private static String customMiniTestingServiceName;
60+
5861
private static void setupRateLimitWatcher (AccountSettings settings) {
5962

6063
scheduler.scheduleAtFixedRate(new Runnable() {
@@ -112,6 +115,18 @@ private static void setTestingRunConfig(TestingRun testingRun, TestingRunResultS
112115
}
113116
}
114117

118+
public static void modifyHybridTestingSettingWithCustomName() {
119+
scheduler.scheduleAtFixedRate(new Runnable() {
120+
public void run() {
121+
customMiniTestingServiceName = System.getenv("MINI_TESTING_NAME");
122+
if(customMiniTestingServiceName == null || customMiniTestingServiceName.trim().isEmpty()) {
123+
customMiniTestingServiceName = "Default_" + UUID.randomUUID().toString().substring(0, 4);
124+
}
125+
dataActor.modifyHybridTestingSettingWithCustomName(RuntimeMode.isHybridDeployment(), customMiniTestingServiceName);
126+
}
127+
}, 0, 5, TimeUnit.MINUTES);
128+
}
129+
115130
public static void main(String[] args) throws InterruptedException {
116131
AccountSettings accountSettings = dataActor.fetchAccountSettings();
117132
setupRateLimitWatcher(accountSettings);
@@ -155,13 +170,13 @@ public void run() {
155170
long startDetailed = System.currentTimeMillis();
156171
int delta = start - 20*60;
157172

158-
TestingRunResultSummary trrs = dataActor.findPendingTestingRunResultSummary(start, delta);
173+
TestingRunResultSummary trrs = dataActor.findPendingTestingRunResultSummary(start, delta, customMiniTestingServiceName);
159174
boolean isSummaryRunning = trrs != null && trrs.getState().equals(State.RUNNING);
160175
TestingRun testingRun;
161176
ObjectId summaryId = null;
162177
if (trrs == null) {
163178
delta = Context.now() - 20*60;
164-
testingRun = dataActor.findPendingTestingRun(delta);
179+
testingRun = dataActor.findPendingTestingRun(delta, customMiniTestingServiceName);
165180
} else {
166181
summaryId = trrs.getId();
167182
loggerMaker.infoAndAddToDb("Found trrs " + trrs.getHexId() + " for account: " + accountId);

libs/dao/src/main/java/com/akto/dto/Account.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package com.akto.dto;
22

3+
import java.util.List;
4+
35
public class Account {
46
private int id;
57
private String name;
@@ -15,6 +17,8 @@ public class Account {
1517
public static final String HYBRID_TESTING_ENABLED = "hybridTestingEnabled";
1618
private boolean hybridTestingEnabled;
1719
private boolean mergingRunning = false;
20+
public static final String MINI_TESTING_HEARTBEAT = "miniTestingHeartbeat";
21+
private List<MiniTestingServiceHeartbeat> miniTestingHeartbeat;
1822

1923
private int mergingInitiateTs = 0;
2024

@@ -104,4 +108,13 @@ public boolean getHybridTestingEnabled() {
104108
public void setHybridTestingEnabled(boolean hybridTestingEnabled) {
105109
this.hybridTestingEnabled = hybridTestingEnabled;
106110
}
111+
112+
public List<MiniTestingServiceHeartbeat> getMiniTestingHeartbeat() {
113+
return miniTestingHeartbeat;
114+
}
115+
116+
public void setMiniTestingHeartbeat(List<MiniTestingServiceHeartbeat> miniTestingHeartbeat) {
117+
this.miniTestingHeartbeat = miniTestingHeartbeat;
118+
}
119+
107120
}
Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
package com.akto.dto;
2+
3+
public class MiniTestingServiceHeartbeat {
4+
5+
public static final String MINI_TESTING_SERVICE_NAME = "miniTestingServiceName";
6+
private String miniTestingServiceName;
7+
8+
public static final String LAST_HEARTBEAT_TIME_STAMP = "lastHeartbeatTimeStamp";
9+
private int lastHeartbeatTimestamp;
10+
11+
public MiniTestingServiceHeartbeat() {}
12+
13+
public MiniTestingServiceHeartbeat(String miniTestingServiceName, int lastHeartbeatTimestamp) {
14+
this.miniTestingServiceName = miniTestingServiceName;
15+
this.lastHeartbeatTimestamp = lastHeartbeatTimestamp;
16+
}
17+
18+
public String getMiniTestingServiceName() {
19+
return miniTestingServiceName;
20+
}
21+
22+
public void setMiniTestingServiceName(String miniTestingServiceName) {
23+
this.miniTestingServiceName = miniTestingServiceName;
24+
}
25+
26+
public int getLastHeartbeatTimestamp() {
27+
return lastHeartbeatTimestamp;
28+
}
29+
30+
public void setLastHeartbeatTimestamp(int lastHeartbeatTimestamp) {
31+
this.lastHeartbeatTimestamp = lastHeartbeatTimestamp;
32+
}
33+
}

libs/dao/src/main/java/com/akto/dto/testing/TestingRun.java

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,9 @@ public class TestingRun {
3737
public static final String NAME = "name";
3838
private String name;
3939

40+
public static final String MINI_TESTING_SERVICE_NAME = "miniTestingServiceName";
41+
private String miniTestingServiceName;
42+
4043
public enum TestingRunType{
4144
ONE_TIME, RECURRING, CI_CD, CONTINUOUS_TESTING
4245
}
@@ -61,7 +64,7 @@ public TestingRun(int scheduleTimestamp, String userEmail, TestingEndpoints test
6164
this.triggeredBy = triggeredBy;
6265
this.sendSlackAlert = sendSlackAlert;
6366
}
64-
public TestingRun(int scheduleTimestamp, String userEmail, TestingEndpoints testingEndpoints, int testIdConfig, State state, int periodInSeconds, String name, int testRunTime, int maxConcurrentRequests, boolean sendSlackAlert) {
67+
public TestingRun(int scheduleTimestamp, String userEmail, TestingEndpoints testingEndpoints, int testIdConfig, State state, int periodInSeconds, String name, int testRunTime, int maxConcurrentRequests, boolean sendSlackAlert, String miniTestingServiceName) {
6568
this.scheduleTimestamp = scheduleTimestamp;
6669
this.testRunTime = testRunTime;
6770
this.maxConcurrentRequests = maxConcurrentRequests;
@@ -74,6 +77,7 @@ public TestingRun(int scheduleTimestamp, String userEmail, TestingEndpoints test
7477
this.periodInSeconds = periodInSeconds;
7578
this.name = name;
7679
this.sendSlackAlert = sendSlackAlert;
80+
this.miniTestingServiceName = miniTestingServiceName;
7781
}
7882

7983
public TestingRunConfig getTestingRunConfig() {
@@ -209,6 +213,14 @@ public void setSendSlackAlert(boolean sendSlackAlert) {
209213
this.sendSlackAlert = sendSlackAlert;
210214
}
211215

216+
public String getMiniTestingServiceName() {
217+
return miniTestingServiceName;
218+
}
219+
220+
public void setMiniTestingServiceName(String miniTestingServiceName) {
221+
this.miniTestingServiceName = miniTestingServiceName;
222+
}
223+
212224
@Override
213225
public String toString() {
214226
return "{" +

libs/utils/src/main/java/com/akto/data_actor/ClientActor.java

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1309,10 +1309,11 @@ public TestingRunResultSummary createTRRSummaryIfAbsent(String testingRunHexId,
13091309
}
13101310
}
13111311

1312-
public TestingRun findPendingTestingRun(int delta) {
1312+
public TestingRun findPendingTestingRun(int delta, String miniTestingServiceName) {
13131313
Map<String, List<String>> headers = buildHeaders();
13141314
BasicDBObject obj = new BasicDBObject();
13151315
obj.put("delta", delta);
1316+
obj.put("miniTestingServiceName", miniTestingServiceName);
13161317
OriginalHttpRequest request = new OriginalHttpRequest(url + "/findPendingTestingRun", "", "POST", obj.toString(), headers, "");
13171318
try {
13181319
OriginalHttpResponse response = ApiExecutor.sendRequest(request, true, null, false, null);
@@ -1364,11 +1365,12 @@ public static <T> T decode(Codec<T> codec, Document doc){
13641365
return codec.decode(bsonReader, DecoderContext.builder().build());
13651366
}
13661367

1367-
public TestingRunResultSummary findPendingTestingRunResultSummary(int now, int delta) {
1368+
public TestingRunResultSummary findPendingTestingRunResultSummary(int now, int delta, String miniTestingServiceName) {
13681369
Map<String, List<String>> headers = buildHeaders();
13691370
BasicDBObject obj = new BasicDBObject();
13701371
obj.put("now", now);
13711372
obj.put("delta", delta);
1373+
obj.put("miniTestingServiceName", miniTestingServiceName);
13721374
OriginalHttpRequest request = new OriginalHttpRequest(url + "/findPendingTestingRunResultSummary", "", "POST", obj.toString(), headers, "");
13731375
try {
13741376
OriginalHttpResponse response = ApiExecutor.sendRequest(request, true, null, false, null);
@@ -3003,6 +3005,25 @@ public void modifyHybridTestingSetting(boolean hybridTestingEnabled) {
30033005
}
30043006
}
30053007

3008+
public void modifyHybridTestingSettingWithCustomName(boolean hybridTestingEnabled, String serviceName) {
3009+
Map<String, List<String>> headers = buildHeaders();
3010+
BasicDBObject obj = new BasicDBObject();
3011+
obj.put("hybridTestingEnabled", hybridTestingEnabled);
3012+
obj.put("miniTestingServiceName", serviceName);
3013+
OriginalHttpRequest request = new OriginalHttpRequest(url + "/modifyTestingSettingWithCustomName", "", "POST", obj.toString(), headers, "");
3014+
try {
3015+
OriginalHttpResponse response = ApiExecutor.sendRequest(request, true, null, false, null);
3016+
String responsePayload = response.getBody();
3017+
if (response.getStatusCode() != 200 || responsePayload == null) {
3018+
loggerMaker.errorAndAddToDb("non 2xx response in modifyTestingSettingWithCustomName", LoggerMaker.LogDb.RUNTIME);
3019+
return;
3020+
}
3021+
} catch (Exception e) {
3022+
loggerMaker.errorAndAddToDb("error in modifyTestingSettingWithCustomName" + e, LoggerMaker.LogDb.RUNTIME);
3023+
return;
3024+
}
3025+
}
3026+
30063027
public void insertTestingLog(Log log) {
30073028
Map<String, List<String>> headers = buildHeaders();
30083029
BasicDBObject obj = new BasicDBObject();

libs/utils/src/main/java/com/akto/data_actor/DataActor.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -110,9 +110,9 @@ public abstract class DataActor {
110110

111111
public abstract TestingRunResultSummary createTRRSummaryIfAbsent(String testingRunHexId, int start);
112112

113-
public abstract TestingRun findPendingTestingRun(int delta);
113+
public abstract TestingRun findPendingTestingRun(int delta, String miniTestingServiceName);
114114

115-
public abstract TestingRunResultSummary findPendingTestingRunResultSummary(int now, int delta);
115+
public abstract TestingRunResultSummary findPendingTestingRunResultSummary(int now, int delta, String miniTestingServiceName);
116116

117117
public abstract TestingRun findTestingRun(String testingRunId);
118118

@@ -228,6 +228,8 @@ public abstract class DataActor {
228228

229229
public abstract void modifyHybridTestingSetting(boolean hybridTestingEnabled);
230230

231+
public abstract void modifyHybridTestingSettingWithCustomName(boolean hybridTestingEnabled, String serviceName);
232+
231233
public abstract void insertTestingLog(Log log);
232234

233235
public abstract EndpointLogicalGroup fetchEndpointLogicalGroupById(String endpointLogicalGroupId);

libs/utils/src/main/java/com/akto/data_actor/DbActor.java

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -242,12 +242,12 @@ public TestingRunResultSummary createTRRSummaryIfAbsent(String testingRunHexId,
242242
return DbLayer.createTRRSummaryIfAbsent(testingRunHexId, start);
243243
}
244244

245-
public TestingRun findPendingTestingRun(int delta) {
246-
return DbLayer.findPendingTestingRun(delta);
245+
public TestingRun findPendingTestingRun(int delta, String miniTestingServiceName) {
246+
return DbLayer.findPendingTestingRun(delta, miniTestingServiceName);
247247
}
248248

249-
public TestingRunResultSummary findPendingTestingRunResultSummary(int now, int delta) {
250-
return DbLayer.findPendingTestingRunResultSummary(now, delta);
249+
public TestingRunResultSummary findPendingTestingRunResultSummary(int now, int delta, String miniTestingServiceName) {
250+
return DbLayer.findPendingTestingRunResultSummary(now, delta, miniTestingServiceName);
251251
}
252252

253253
public TestingRunConfig findTestingRunConfig(int testIdConfig) {
@@ -482,6 +482,10 @@ public void modifyHybridTestingSetting(boolean hybridTestingEnabled) {
482482
DbLayer.modifyHybridTestingSetting(hybridTestingEnabled);
483483
}
484484

485+
public void modifyHybridTestingSettingWithCustomName(boolean hybridTestingEnabled, String serviceName) {
486+
DbLayer.modifyHybridTestingSettingWithCustomName(hybridTestingEnabled, serviceName);
487+
}
488+
485489
public void insertTestingLog(Log log) {
486490
DbLayer.insertTestingLog(log);
487491
}

0 commit comments

Comments
 (0)