Эх сурвалжийг харах

feat: ChatService/QuotaService逻辑更新+测试

Ultraworked with Sisyphus
liaoxg 3 сар өмнө
parent
commit
96ceafaead

+ 6 - 9
num-server/src/main/java/com/etotem/num/service/ChatService.java

@@ -13,14 +13,14 @@ public class ChatService {
     private final ChatMessageRepository chatMessageRepository;
     private final ChartService chartService;
     private final DifyService difyService;
-    private final UserService userService;
+    private final QuotaService quotaService;
 
     public ChatService(ChatMessageRepository chatMessageRepository, ChartService chartService,
-                       DifyService difyService, UserService userService) {
+                       DifyService difyService, QuotaService quotaService) {
         this.chatMessageRepository = chatMessageRepository;
         this.chartService = chartService;
         this.difyService = difyService;
-        this.userService = userService;
+        this.quotaService = quotaService;
     }
 
     /**
@@ -28,15 +28,15 @@ public class ChatService {
      * Checks quota, saves messages, calls Dify, returns response.
      */
     public String sendMessage(Long userId, Long chartRecordId, String query, String audioUrl) {
-        // Check daily chat quota
-        userService.checkDailyQuota(userId, "chat");
-
         // Verify user owns this chart record
         ChartRecord record = chartService.getById(chartRecordId);
         if (!record.getUserId().equals(userId)) {
             throw new com.etotem.num.common.BizException(1006, "无权访问此能量盘");
         }
 
+        // Check and consume consultation quota (person + daily limits)
+        quotaService.checkAndConsume(userId, record.getName());
+
         // Save user message
         ChatMessage userMsg = new ChatMessage();
         userMsg.setChartRecordId(chartRecordId);
@@ -63,9 +63,6 @@ public class ChatService {
         aiMsg.setContent(response);
         chatMessageRepository.save(aiMsg);
 
-        // Consume chat quota
-        userService.consumeQuota(userId, "chat");
-
         return response;
     }
 

+ 1 - 1
num-server/src/main/java/com/etotem/num/service/QuotaService.java

@@ -56,7 +56,7 @@ public class QuotaService {
         Optional<UserConsultationQuota> existing =
                 quotaRepository.findByUserIdAndPersonName(userId, personName);
 
-        if (existing.isEmpty()) {
+        if (!existing.isPresent()) {
             // --- New person: enforce person limit ---
             if (personLimit > 0) {
                 long currentPersonCount = quotaRepository.countByUserId(userId);

+ 19 - 10
num-server/src/test/java/com/etotem/num/service/ChatServiceTest.java

@@ -1,5 +1,6 @@
 package com.etotem.num.service;
 
+import com.etotem.num.common.BizException;
 import com.etotem.num.common.BizException;
 import com.etotem.num.entity.ChatMessage;
 import com.etotem.num.entity.ChartRecord;
@@ -34,7 +35,7 @@ class ChatServiceTest {
     private DifyService difyService;
 
     @MockBean
-    private UserService userService;
+    private QuotaService quotaService;
 
     // ==================== sendMessage ====================
 
@@ -43,6 +44,7 @@ class ChatServiceTest {
         ChartRecord record = new ChartRecord();
         record.setId(10L);
         record.setUserId(1L);
+        record.setName("测试用户");
         record.setChartData("{\"numbers\": [1,2,3]}");
 
         when(chartService.getById(10L)).thenReturn(record);
@@ -63,15 +65,20 @@ class ChatServiceTest {
                 "ai".equals(msg.getRole()) && "AI回复内容".equals(msg.getContent())
                         && 10L == msg.getChartRecordId()));
 
-        // Quota checked and consumed
-        verify(userService).checkDailyQuota(1L, "chat");
-        verify(userService).consumeQuota(1L, "chat");
+        // Quota checked and consumed via QuotaService
+        verify(quotaService).checkAndConsume(1L, "测试用户");
     }
 
     @Test
     void testSendMessage_quotaExceeded() {
-        doThrow(new BizException(1003, "超过每日AI解读次数"))
-                .when(userService).checkDailyQuota(1L, "chat");
+        ChartRecord record = new ChartRecord();
+        record.setId(10L);
+        record.setUserId(1L);
+        record.setName("测试用户");
+
+        when(chartService.getById(10L)).thenReturn(record);
+        doThrow(new BizException(403, "已达咨询人数上限,最多可咨询3人"))
+                .when(quotaService).checkAndConsume(1L, "测试用户");
 
         assertThrows(BizException.class, () -> chatService.sendMessage(1L, 10L, "test", null));
         verify(chatMessageRepository, never()).save(any());
@@ -98,6 +105,7 @@ class ChatServiceTest {
         ChartRecord record = new ChartRecord();
         record.setId(10L);
         record.setUserId(1L);
+        record.setName("测试用户");
         record.setChartData("{}");
 
         when(chartService.getById(10L)).thenReturn(record);
@@ -116,10 +124,11 @@ class ChatServiceTest {
     }
 
     @Test
-    void testSendMessage_quotaConsumedAfterDifyCall() {
+    void testSendMessage_quotaCheckedBeforeDifyCall() {
         ChartRecord record = new ChartRecord();
         record.setId(10L);
         record.setUserId(1L);
+        record.setName("测试用户");
         record.setChartData("{}");
 
         when(chartService.getById(10L)).thenReturn(record);
@@ -129,10 +138,10 @@ class ChatServiceTest {
 
         chatService.sendMessage(1L, 10L, "hi", null);
 
-        // Verify consumeQuota is called AFTER invokeChatflow
-        InOrder inOrder = inOrder(userService, difyService);
+        // Verify checkAndConsume is called BEFORE invokeChatflow
+        InOrder inOrder = inOrder(quotaService, difyService);
+        inOrder.verify(quotaService).checkAndConsume(eq(1L), anyString());
         inOrder.verify(difyService).invokeChatflow(anyString(), anyString(), anyString());
-        inOrder.verify(userService).consumeQuota(1L, "chat");
     }
 
     // ==================== getHistory ====================