generate 添加试用用户判断

src/main/java/com/ai/da/service/GenerateService.java

@@ -7,6 +7,7 @@ import com.ai.da.model.dto.GenerateThroughImageTextDTO;
 import com.ai.da.model.vo.GenerateCaptionVO;
 import com.ai.da.model.vo.GenerateCollectionVO;
 import com.ai.da.model.vo.GenerateLikeVO;
+import com.ai.da.model.vo.PrepareForGenerateVO;
 import com.baomidou.mybatisplus.extension.service.IService;
 
 import java.util.List;
@@ -27,7 +28,7 @@ public interface GenerateService extends IService<Generate> {
 
     GenerateCollectionVO getGenerateResult(String uniqueId);
 
-    String prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO);
+    PrepareForGenerateVO prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO);
 
     Long getRankPosition(String uniqueId);
 

src/main/java/com/ai/da/service/impl/GenerateServiceImpl.java

@@ -314,7 +314,7 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
     }
 
     @Override
-    public String prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO) {
+    public PrepareForGenerateVO prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO) {
         // 1、参数检查，判断必须参数是否为空
         if (Objects.isNull(generateThroughImageTextDTO.getUserId())) {
             throw new BusinessException("userId cannot be empty");
@@ -323,6 +323,16 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
         if (!GenerateModeEnum.getGenerateModeList().contains(generateType)) {
             throw new BusinessException("unknown.generate.type");
         }
+
+        // 判断试用用户是否还有剩余试用机会
+        int trialsCount = 0;
+        if (generateThroughImageTextDTO.getIsTestUser()){
+            trialsCount = getTrialsCount(generateThroughImageTextDTO.getUserId(), generateThroughImageTextDTO.getLevel1Type());
+            if (trialsCount >= 2){
+                return new PrepareForGenerateVO(0);
+            }
+        }
+
         String text = generateThroughImageTextDTO.getText();
         Long elementId = generateThroughImageTextDTO.getCollectionElementId();
         validateGeneraType(new Generate(), text, elementId, generateType);
@@ -361,7 +371,7 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
         rabbitMQService.publishMessage(jsonString);
 
         // 5、返回唯一id
-        return uuid;
+        return new PrepareForGenerateVO(uuid, 2 - trialsCount);
     }
 
     @Override
@@ -458,4 +468,28 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
         GenerateCancel generateCancel = new GenerateCancel(userId, uniqueId, DateUtil.getByTimeZone(timeZone));
         generateCancelMapper.insert(generateCancel);
     }
+
+    // 判断试用用户试用generate机会是否使用完毕
+    private int getTrialsCount(Long userId, String level1Type){
+        List<Generate> getGenerateList = getGenerateByAccountId(userId, level1Type);
+        int trialsCount ;
+        if (getGenerateList.isEmpty()){
+            trialsCount = 0;
+        } else if (getGenerateList.size() == 1) {
+            trialsCount = 1;
+        } else if (getGenerateList.size() == 2) {
+            trialsCount = 2;
+        }else {
+            trialsCount = 2;
+        }
+        return trialsCount;
+    }
+
+    public List<Generate> getGenerateByAccountId(Long accountId, String level1Type){
+        QueryWrapper<Generate> qw = new QueryWrapper<>();
+        qw.eq("account_id",accountId);
+        qw.eq("level1_type", level1Type);
+
+        return baseMapper.selectList(qw);
+    }
 }