Merge pull request #2066 from akto-api-security/feature/custom_mini_testing_selection_cyborg

shivam-rawat-akto · web-flow · commit 43ec1667e52a · 2025-05-01T01:43:40.000+05:30
feat: custom mini testing module selection for cyborg
diff --git a/apps/dashboard/src/main/java/com/akto/action/testing/StartTestAction.java b/apps/dashboard/src/main/java/com/akto/action/testing/StartTestAction.java
@@ -159,7 +159,7 @@ private TestingRun createTestingRun(int scheduleTimestamp, int periodInSeconds)
 
         return new TestingRun(scheduleTimestamp, user.getLogin(),
                 testingEndpoints, testIdConfig, State.SCHEDULED, periodInSeconds, testName, this.testRunTime,
-                this.maxConcurrentRequests, this.sendSlackAlert);
+                this.maxConcurrentRequests, this.sendSlackAlert, null);
     }
 
     private List<String> selectedTests;
diff --git a/apps/database-abstractor/src/main/java/com/akto/action/DbAction.java b/apps/database-abstractor/src/main/java/com/akto/action/DbAction.java
@@ -2347,7 +2347,7 @@ public String findTestSubCategoriesByTestSuiteId() {
         }
         return Action.SUCCESS.toUpperCase();
     }
-    
+
     public String getCurrentTestingRunDetailsFromEditor(){
         try {
             testingRunPlayground = DbLayer.getCurrentTestingRunDetailsFromEditor(this.ts);
@@ -3485,7 +3485,7 @@ public void setTestSuiteId(List<String> testSuiteId) {
     public List<String> getTestSuiteTestSubCategories() {
         return testSuiteTestSubCategories;
     }
-    
+
     public TestingRunPlayground getTestingRunPlayground() {
         return testingRunPlayground;
     }
diff --git a/libs/dao/src/main/java/com/akto/dto/monitoring/ModuleInfo.java b/libs/dao/src/main/java/com/akto/dto/monitoring/ModuleInfo.java
@@ -13,6 +13,7 @@ public class ModuleInfo {
     private int startedTs;
     public static final String LAST_HEARTBEAT_RECEIVED = "lastHeartbeatReceived";
     private int lastHeartbeatReceived;
+    public static final String NAME = "name";
     private String name;
 
     public ModuleType getModuleType() {
diff --git a/libs/dao/src/main/java/com/akto/dto/testing/TestingRun.java b/libs/dao/src/main/java/com/akto/dto/testing/TestingRun.java
@@ -37,6 +37,9 @@ public class TestingRun {
     public static final String NAME = "name";
     private String name;
 
+    public static final String MINI_TESTING_SERVICE_NAME = "miniTestingServiceName";
+    private String miniTestingServiceName;
+
     public enum TestingRunType{
         ONE_TIME, RECURRING, CI_CD, CONTINUOUS_TESTING
     }
@@ -61,7 +64,7 @@ public TestingRun(int scheduleTimestamp, String userEmail, TestingEndpoints test
         this.triggeredBy = triggeredBy;
         this.sendSlackAlert = sendSlackAlert;
     }
-    public TestingRun(int scheduleTimestamp, String userEmail, TestingEndpoints testingEndpoints, int testIdConfig, State state, int periodInSeconds, String name, int testRunTime, int maxConcurrentRequests, boolean sendSlackAlert) {
+    public TestingRun(int scheduleTimestamp, String userEmail, TestingEndpoints testingEndpoints, int testIdConfig, State state, int periodInSeconds, String name, int testRunTime, int maxConcurrentRequests, boolean sendSlackAlert, String miniTestingServiceName) {
         this.scheduleTimestamp = scheduleTimestamp;
         this.testRunTime = testRunTime;
         this.maxConcurrentRequests = maxConcurrentRequests;
@@ -74,6 +77,7 @@ public TestingRun(int scheduleTimestamp, String userEmail, TestingEndpoints test
         this.periodInSeconds = periodInSeconds;
         this.name = name;
         this.sendSlackAlert = sendSlackAlert;
+        this.miniTestingServiceName = miniTestingServiceName;
     }
 
     public TestingRunConfig getTestingRunConfig() {
@@ -209,6 +213,14 @@ public void setSendSlackAlert(boolean sendSlackAlert) {
         this.sendSlackAlert = sendSlackAlert;
     }
 
+    public String getMiniTestingServiceName() {
+        return miniTestingServiceName;
+    }
+
+    public void setMiniTestingServiceName(String miniTestingServiceName) {
+        this.miniTestingServiceName = miniTestingServiceName;
+    }
+
     @Override
     public String toString() {
         return "{" +
diff --git a/libs/utils/src/main/java/com/akto/data_actor/DbLayer.java b/libs/utils/src/main/java/com/akto/data_actor/DbLayer.java
@@ -11,6 +11,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 
 import com.akto.bulk_update_util.ApiInfoBulkUpdate;
 import com.akto.dao.*;
@@ -106,6 +107,12 @@ public class DbLayer {
 
     private static final LoggerMaker loggerMaker = new LoggerMaker(DbLayer.class, LoggerMaker.LogDb.DASHBOARD);
 
+    private static final ConcurrentHashMap<Integer, Integer> lastUpdatedTsMap = new ConcurrentHashMap<>();
+
+    private static int getLastUpdatedTsForAccount(int accountId) {
+        return lastUpdatedTsMap.computeIfAbsent(accountId, k -> 0);
+    }
+
     public DbLayer() {
     }
 
@@ -138,8 +145,60 @@ public static void updateModuleInfo(ModuleInfo moduleInfo) {
                         Updates.setOnInsert(ModuleInfo.MODULE_TYPE, moduleInfo.getModuleType()),
                         Updates.setOnInsert(ModuleInfo.STARTED_TS, moduleInfo.getStartedTs()),
                         Updates.setOnInsert(ModuleInfo.CURRENT_VERSION, moduleInfo.getCurrentVersion()),
+                        Updates.setOnInsert(ModuleInfo.NAME, moduleInfo.getName()),
                         Updates.set(ModuleInfo.LAST_HEARTBEAT_RECEIVED, moduleInfo.getLastHeartbeatReceived())
                 ), updateOptions);
+        if (moduleInfo.getModuleType() == ModuleInfo.ModuleType.MINI_TESTING) {
+            //Only for mini-testing heartbeat mark testing run failed state
+            if (Context.now() - getLastUpdatedTsForAccount(Context.accountId.get()) > 10 * 60) {
+                try {
+                    fetchAndFailOutdatedTests();
+                } catch (Exception e) {
+                    loggerMaker.errorAndAddToDb("Error while updating out");
+                    //Ignore
+                }
+                lastUpdatedTsMap.put(Context.accountId.get(), Context.now());
+            }
+        }
+    }
+
+    public static void fetchAndFailOutdatedTests() {
+        List<TestingRun> testingRunList = TestingRunDao.instance.findAll(
+                Filters.and(
+                        Filters.or(
+                                Filters.eq(TestingRun.STATE, State.SCHEDULED),
+                                Filters.eq(TestingRun.STATE, State.RUNNING)
+                        ),
+                        Filters.exists(ModuleInfo.NAME, true),
+                        Filters.ne(ModuleInfo.NAME, null)
+                )
+        );
+
+        for (TestingRun testingRun : testingRunList) {
+            String miniTestingServiceName = testingRun.getMiniTestingServiceName();
+            if(!miniTestingServiceName.isEmpty()) {
+                List<ModuleInfo> moduleInfos = ModuleInfoDao.instance.
+                        findAll(Filters.eq(ModuleInfo.NAME, miniTestingServiceName), 0, 1,
+                                Sorts.descending(ModuleInfo.LAST_HEARTBEAT_RECEIVED));
+                boolean isValid = true;
+                if (moduleInfos != null && !moduleInfos.isEmpty()) {
+                    isValid = Context.now() - moduleInfos.get(0).getLastHeartbeatReceived() < 20 * 60;
+                }
+
+                if(!isValid) {
+                    Bson filter = Filters.or(
+                            Filters.eq(TestingRun.STATE, State.SCHEDULED),
+                            Filters.eq(TestingRun.STATE, State.RUNNING));
+                    TestingRunDao.instance.updateOne(
+                            Filters.and(filter, Filters.eq(Constants.ID, testingRun.getId())),
+                            Updates.set(TestingRun.STATE, State.FAILED));
+                    TestingRunResultSummariesDao.instance.updateOneNoUpsert(
+                            Filters.eq(TestingRunResultSummary.TESTING_RUN_ID, testingRun.getId()),
+                            Updates.set(TestingRunResultSummary.STATE, State.FAILED)
+                    );
+                }
+            }
+        }
     }
 
 

Original file line number	Diff line number	Diff line change
`@@ -159,7 +159,7 @@ private TestingRun createTestingRun(int scheduleTimestamp, int periodInSeconds)`
`159`	`159`
`160`	`160`	`return new TestingRun(scheduleTimestamp, user.getLogin(),`
`161`	`161`	`testingEndpoints, testIdConfig, State.SCHEDULED, periodInSeconds, testName, this.testRunTime,`
`162`		`- this.maxConcurrentRequests, this.sendSlackAlert);`
	`162`	`+ this.maxConcurrentRequests, this.sendSlackAlert, null);`
`163`	`163`	`}`
`164`	`164`
`165`	`165`	`private List<String> selectedTests;`
Original file line number	Diff line number	Diff line change
`@@ -2347,7 +2347,7 @@ public String findTestSubCategoriesByTestSuiteId() {`
`2347`	`2347`	`}`
`2348`	`2348`	`return Action.SUCCESS.toUpperCase();`
`2349`	`2349`	`}`
`2350`		`-`
	`2350`	`+`
`2351`	`2351`	`public String getCurrentTestingRunDetailsFromEditor(){`
`2352`	`2352`	`try {`
`2353`	`2353`	`testingRunPlayground = DbLayer.getCurrentTestingRunDetailsFromEditor(this.ts);`
`@@ -3485,7 +3485,7 @@ public void setTestSuiteId(List<String> testSuiteId) {`
`3485`	`3485`	`public List<String> getTestSuiteTestSubCategories() {`
`3486`	`3486`	`return testSuiteTestSubCategories;`
`3487`	`3487`	`}`
`3488`		`-`
	`3488`	`+`
`3489`	`3489`	`public TestingRunPlayground getTestingRunPlayground() {`
`3490`	`3490`	`return testingRunPlayground;`
`3491`	`3491`	`}`