From abe9409ecbe5ac752dd6f14100733d1b3739ac09 Mon Sep 17 00:00:00 2001
From: Junjie <fallin.jie@qq.com>
Date: 星期一, 27 四月 2026 12:14:36 +0800
Subject: [PATCH] test: cover auto tune agent safety paths

---
 src/main/java/com/zy/ai/mcp/tool/AutoTuneMcpTools.java |   23 ++++++++++++++++++-----
 1 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/src/main/java/com/zy/ai/mcp/tool/AutoTuneMcpTools.java b/src/main/java/com/zy/ai/mcp/tool/AutoTuneMcpTools.java
index e822cf3..dcbe782 100644
--- a/src/main/java/com/zy/ai/mcp/tool/AutoTuneMcpTools.java
+++ b/src/main/java/com/zy/ai/mcp/tool/AutoTuneMcpTools.java
@@ -26,6 +26,7 @@
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.function.LongSupplier;
 
 @Component
 @RequiredArgsConstructor
@@ -40,6 +41,7 @@
     private final AiAutoTuneJobService aiAutoTuneJobService;
     private final AiAutoTuneChangeService aiAutoTuneChangeService;
     private final ConcurrentMap<String, DryRunPreview> dryRunPreviews = new ConcurrentHashMap<>();
+    private LongSupplier currentTimeMillisSupplier = System::currentTimeMillis;
 
     @Tool(name = "dispatch_get_auto_tune_snapshot", description = "鑾峰彇WCS鑷姩璋冨弬鎵�闇�鐨勮皟搴﹀揩鐓с�佺珯鐐硅繍琛屾�併�佹嫇鎵戝閲忓拰褰撳墠鍙啓鍙傛暟")
     public AutoTuneSnapshot getAutoTuneSnapshot() {
@@ -164,7 +166,7 @@
         if (preview == null) {
             throw new IllegalArgumentException("dryRunToken is missing, expired, or already used.");
         }
-        if (preview.isExpired()) {
+        if (preview.isExpired(currentTimeMillis())) {
             throw new IllegalArgumentException("dryRunToken is expired. Run dryRun=true again.");
         }
         if (!preview.getFingerprint().equals(fingerprint)) {
@@ -175,13 +177,14 @@
     private String createDryRunToken(String fingerprint) {
         cleanExpiredDryRunPreviews();
         String token = UUID.randomUUID().toString();
-        dryRunPreviews.put(token, new DryRunPreview(fingerprint, System.currentTimeMillis() + DRY_RUN_TOKEN_TTL_MILLIS));
+        dryRunPreviews.put(token, new DryRunPreview(fingerprint, currentTimeMillis() + DRY_RUN_TOKEN_TTL_MILLIS));
         return token;
     }
 
     private void cleanExpiredDryRunPreviews() {
+        long currentTimeMillis = currentTimeMillis();
         for (Map.Entry<String, DryRunPreview> entry : dryRunPreviews.entrySet()) {
-            if (entry.getValue() == null || entry.getValue().isExpired()) {
+            if (entry.getValue() == null || entry.getValue().isExpired(currentTimeMillis)) {
                 dryRunPreviews.remove(entry.getKey());
             }
         }
@@ -228,6 +231,16 @@
         return value == null || value.trim().isEmpty();
     }
 
+    private long currentTimeMillis() {
+        return currentTimeMillisSupplier.getAsLong();
+    }
+
+    void setCurrentTimeMillisSupplier(LongSupplier currentTimeMillisSupplier) {
+        this.currentTimeMillisSupplier = currentTimeMillisSupplier == null
+                ? System::currentTimeMillis
+                : currentTimeMillisSupplier;
+    }
+
     private static class DryRunPreview {
         private final String fingerprint;
         private final long expireAtMillis;
@@ -241,8 +254,8 @@
             return fingerprint;
         }
 
-        boolean isExpired() {
-            return System.currentTimeMillis() > expireAtMillis;
+        boolean isExpired(long currentTimeMillis) {
+            return currentTimeMillis > expireAtMillis;
         }
     }
 }

--
Gitblit v1.9.1