Add chargebacks and refunds (#25)

src/main/antlr4/com.rbkmoney.fraudo/Fraudo.g4

@@ -17,9 +17,13 @@ expression
  | count                                          #countExpression
  | count_success                                  #countSuccessExpression
  | count_error                                    #countErrorExpression
+ | count_chargeback                               #countChargebackExpression
+ | count_refund                                   #countRefundExpression
  | sum                                            #sumExpression
  | sum_success                                    #sumSuccessExpression
  | sum_error                                      #sumErrorExpression
+ | sum_chargeback                                 #sumChargebackExpression
+ | sum_refund                                     #sumRefundExpression
  | unique                                         #uniqueExpression
  | in                                             #inFunctionExpression
  | in_white_list                                  #inWhiteListExpression
@@ -68,6 +72,14 @@ count_error
  : 'countError' LPAREN STRING time_window DELIMETER STRING (group_by)? RPAREN
  ;
 
+count_chargeback
+ : 'countChargeback' LPAREN STRING time_window (group_by)? RPAREN
+ ;
+
+count_refund
+ : 'countRefund' LPAREN STRING time_window (group_by)? RPAREN
+ ;
+
 sum
  : 'sum' LPAREN STRING time_window (group_by)? RPAREN
  ;
@@ -80,6 +92,14 @@ sum_error
  : 'sumError' LPAREN STRING time_window DELIMETER STRING (group_by)? RPAREN
  ;
 
+sum_chargeback
+ : 'sumChargeback' LPAREN STRING time_window (group_by)? RPAREN
+ ;
+
+sum_refund
+ : 'sumRefund' LPAREN STRING time_window (group_by)? RPAREN
+ ;
+
 unique
  : 'unique' LPAREN STRING DELIMETER STRING time_window (group_by)? RPAREN
  ;

src/main/java/com/rbkmoney/fraudo/aggregator/CountAggregator.java

@@ -12,4 +12,8 @@ public interface CountAggregator<T, U> {
 
     Integer countError(U checkedField, T model, TimeWindow timeWindow, String errorCode, List<U> fields);
 
+    Integer countChargeback(U checkedField, T model, TimeWindow timeWindow, List<U> fields);
+
+    Integer countRefund(U checkedField, T model, TimeWindow timeWindow, List<U> fields);
+
 }

src/main/java/com/rbkmoney/fraudo/aggregator/SumAggregator.java

@@ -12,4 +12,8 @@ public interface SumAggregator<T, U> {
 
     Double sumError(U checkedField, T model, TimeWindow timeWindow, String errorCode, List<U> fields);
 
+    Double sumChargeback(U checkedField, T model, TimeWindow timeWindow, List<U> fields);
+
+    Double sumRefund(U checkedField, T model, TimeWindow timeWindow, List<U> fields);
+
 }

src/main/java/com/rbkmoney/fraudo/utils/key/generator/CountKeyGenerator.java

@@ -38,4 +38,22 @@ public class CountKeyGenerator {
                 resolve);
     }
 
+    public static <T> String generateChargebackKey(ParserRuleContext context, Function<String, T> resolve) {
+        FraudoParser.Count_chargebackContext ctx = (FraudoParser.Count_chargebackContext) context;
+        return CommonKeyGenerator.generateKeyGroupedFunction(ctx.STRING(),
+                ctx.children.get(0),
+                ctx.time_window(),
+                ctx.group_by(),
+                resolve);
+    }
+
+    public static <T> String generateRefundKey(ParserRuleContext context, Function<String, T> resolve) {
+        FraudoParser.Count_refundContext ctx = (FraudoParser.Count_refundContext) context;
+        return CommonKeyGenerator.generateKeyGroupedFunction(ctx.STRING(),
+                ctx.children.get(0),
+                ctx.time_window(),
+                ctx.group_by(),
+                resolve);
+    }
+
 }

src/main/java/com/rbkmoney/fraudo/utils/key/generator/SumKeyGenerator.java

@@ -38,4 +38,22 @@ public class SumKeyGenerator {
                 resolve);
     }
 
+    public static <T> String generateChargebackKey(ParserRuleContext context, Function<String, T> resolve) {
+        FraudoParser.Sum_chargebackContext ctx = (FraudoParser.Sum_chargebackContext) context;
+        return CommonKeyGenerator.generateKeyGroupedFunction(ctx.STRING(),
+                ctx.children.get(0),
+                ctx.time_window(),
+                ctx.group_by(),
+                resolve);
+    }
+
+    public static <T> String generateRefundKey(ParserRuleContext context, Function<String, T> resolve) {
+        FraudoParser.Sum_refundContext ctx = (FraudoParser.Sum_refundContext) context;
+        return CommonKeyGenerator.generateKeyGroupedFunction(ctx.STRING(),
+                ctx.children.get(0),
+                ctx.time_window(),
+                ctx.group_by(),
+                resolve);
+    }
+
 }

src/main/java/com/rbkmoney/fraudo/visitor/CountVisitor.java

@@ -10,4 +10,8 @@ public interface CountVisitor<T> {
 
     Integer visitCountError(FraudoParser.Count_errorContext ctx, T model);
 
+    Integer visitCountChargeback(FraudoParser.Count_chargebackContext ctx, T model);
+
+    Integer visitCountRefund(FraudoParser.Count_refundContext ctx, T model);
+
 }

src/main/java/com/rbkmoney/fraudo/visitor/SumVisitor.java

@@ -10,4 +10,8 @@ public interface SumVisitor<T> {
 
     Double visitSumError(FraudoParser.Sum_errorContext ctx, T model);
 
+    Double visitSumChargeback(FraudoParser.Sum_chargebackContext ctx, T model);
+
+    Double visitSumRefund(FraudoParser.Sum_refundContext ctx, T model);
+
 }

src/main/java/com/rbkmoney/fraudo/visitor/impl/CountVisitorImpl.java

@@ -51,4 +51,26 @@ public class CountVisitorImpl<T, U> implements CountVisitor<T> {
         );
     }
 
+    @Override
+    public Integer visitCountChargeback(FraudoParser.Count_chargebackContext ctx, T model) {
+        String countTarget = TextUtil.safeGetText(ctx.STRING());
+        return countAggregator.countChargeback(
+                fieldResolver.resolveName(countTarget),
+                model,
+                TimeWindowResolver.resolve(ctx.time_window()),
+                groupFieldsResolver.resolve(ctx.group_by())
+        );
+    }
+
+    @Override
+    public Integer visitCountRefund(FraudoParser.Count_refundContext ctx, T model) {
+        String countTarget = TextUtil.safeGetText(ctx.STRING());
+        return countAggregator.countRefund(
+                fieldResolver.resolveName(countTarget),
+                model,
+                TimeWindowResolver.resolve(ctx.time_window()),
+                groupFieldsResolver.resolve(ctx.group_by())
+        );
+    }
+
 }

src/main/java/com/rbkmoney/fraudo/visitor/impl/FirstFindVisitorImpl.java

@@ -174,6 +174,24 @@ public class FirstFindVisitorImpl<T extends BaseModel, U> extends FraudoBaseVisi
         );
     }
 
+    @Override
+    public Object visitCount_chargeback(FraudoParser.Count_chargebackContext ctx) {
+        String key = CountKeyGenerator.generateChargebackKey(ctx, fieldResolver::resolveName);
+        return localFuncCache.get().computeIfAbsent(
+                key,
+                s -> Double.valueOf(countVisitor.visitCountChargeback(ctx, threadLocalModel.get()))
+        );
+    }
+
+    @Override
+    public Object visitCount_refund(FraudoParser.Count_refundContext ctx) {
+        String key = CountKeyGenerator.generateRefundKey(ctx, fieldResolver::resolveName);
+        return localFuncCache.get().computeIfAbsent(
+                key,
+                s -> Double.valueOf(countVisitor.visitCountRefund(ctx, threadLocalModel.get()))
+        );
+    }
+
     @Override
     public Object visitSum(FraudoParser.SumContext ctx) {
         String key = SumKeyGenerator.generate(ctx, fieldResolver::resolveName);
@@ -201,6 +219,24 @@ public class FirstFindVisitorImpl<T extends BaseModel, U> extends FraudoBaseVisi
         );
     }
 
+    @Override
+    public Object visitSum_chargeback(FraudoParser.Sum_chargebackContext ctx) {
+        String key = SumKeyGenerator.generateChargebackKey(ctx, fieldResolver::resolveName);
+        return localFuncCache.get().computeIfAbsent(
+                key,
+                s -> sumVisitor.visitSumChargeback(ctx, threadLocalModel.get())
+        );
+    }
+
+    @Override
+    public Object visitSum_refund(FraudoParser.Sum_refundContext ctx) {
+        String key = SumKeyGenerator.generateRefundKey(ctx, fieldResolver::resolveName);
+        return localFuncCache.get().computeIfAbsent(
+                key,
+                s -> sumVisitor.visitSumRefund(ctx, threadLocalModel.get())
+        );
+    }
+
     @Override
     public Object visitCountry_by(FraudoParser.Country_byContext ctx) {
         String key = CountryKeyGenerator.generate(ctx);

src/main/java/com/rbkmoney/fraudo/visitor/impl/SumVisitorImpl.java

@@ -49,4 +49,26 @@ public class SumVisitorImpl<T, U> implements SumVisitor<T> {
                 groupByModelResolver.resolve(ctx.group_by()));
     }
 
+    @Override
+    public Double visitSumChargeback(FraudoParser.Sum_chargebackContext ctx, T model) {
+        String countTarget = TextUtil.safeGetText(ctx.STRING());
+        return sumAggregator.sumChargeback(
+                fieldResolver.resolveName(countTarget),
+                model,
+                TimeWindowResolver.resolve(ctx.time_window()),
+                groupByModelResolver.resolve(ctx.group_by())
+        );
+    }
+
+    @Override
+    public Double visitSumRefund(FraudoParser.Sum_refundContext ctx, T model) {
+        String countTarget = TextUtil.safeGetText(ctx.STRING());
+        return sumAggregator.sumRefund(
+                fieldResolver.resolveName(countTarget),
+                model,
+                TimeWindowResolver.resolve(ctx.time_window()),
+                groupByModelResolver.resolve(ctx.group_by())
+        );
+    }
+
 }

src/test/java/com/rbkmoney/fraudo/CountTest.java

@@ -25,6 +25,8 @@ public class CountTest extends AbstractPaymentTest {
         when(countAggregator.count(anyObject(), any(), any(), any())).thenReturn(10);
         when(countAggregator.countError(anyObject(), any(), any(), anyString(), any())).thenReturn(6);
         when(countAggregator.countSuccess(anyObject(), any(), any(), any())).thenReturn(4);
+        when(countAggregator.countChargeback(anyObject(), any(), any(), any())).thenReturn(0);
+        when(countAggregator.countRefund(anyObject(), any(), any(), any())).thenReturn(1);
         com.rbkmoney.fraudo.FraudoParser.ParseContext parseContext = getParseContext(resourceAsStream);
         ResultModel result = invokeParse(parseContext);
         assertEquals(ResultStatus.DECLINE, result.getResultStatus());
@@ -38,6 +40,17 @@ public class CountTest extends AbstractPaymentTest {
         assertEquals(ResultStatus.DECLINE, result.getResultStatus());
     }
 
+    @Test
+    public void countCargeRefundTest() throws Exception {
+        InputStream resourceAsStream = CountTest.class.getResourceAsStream("/rules/count_chargeback_refund.frd");
+        when(countAggregator.countChargeback(anyObject(), any(), any(), any())).thenReturn(3);
+        when(countAggregator.countRefund(anyObject(), any(), any(), any())).thenReturn(5);
+        com.rbkmoney.fraudo.FraudoParser.ParseContext parseContext = getParseContext(resourceAsStream);
+        ResultModel result = invokeParse(parseContext);
+        assertEquals(ResultStatus.DECLINE, result.getResultStatus());
+        assertEquals("1", result.getRuleChecked());
+    }
+
     @Test
     public void countGroupByTest() throws Exception {
         InputStream resourceAsStream = CountTest.class.getResourceAsStream("/rules/countGroupBy.frd");

src/test/java/com/rbkmoney/fraudo/SumTest.java

@@ -39,6 +39,16 @@ public class SumTest extends AbstractPaymentTest {
         Assert.assertEquals(0, result.getNotificationsRule().size());
     }
 
+    @Test
+    public void sumChargeRefundTest() throws Exception {
+        InputStream resourceAsStream = SumTest.class.getResourceAsStream("/rules/sum_chargeback_refund.frd");
+        Mockito.when(sumAggregator.sumChargeback(anyObject(), any(), any(), any())).thenReturn(10000.60);
+        Mockito.when(sumAggregator.sumRefund(anyObject(), any(), any(), any())).thenReturn(10000.60);
+        com.rbkmoney.fraudo.FraudoParser.ParseContext parseContext = getParseContext(resourceAsStream);
+        ResultModel result = invokeParse(parseContext);
+        Assert.assertEquals(ResultStatus.ACCEPT, result.getResultStatus());
+    }
+
     @Test
     public void sumGroupByTest() throws Exception {
         InputStream resourceAsStream = SumTest.class.getResourceAsStream("/rules/sumGroupBy.frd");

src/test/resources/rules/amount.frd

@@ -1,2 +1,2 @@
-rule: amount() < 100 AND currency() == "RUB"
+rule: amount() < 100 AND currency() = "RUB"
 -> accept;

src/test/resources/rules/count_chargeback_refund.frd

@@ -0,0 +1,2 @@
+rule: countChargeback("ip", 1444) > 3 OR countRefund("ip", 1444) > 1
+-> decline;

src/test/resources/rules/sum_chargeback_refund.frd

@@ -0,0 +1,2 @@
+rule: sumChargeback("ip", 1444) < 10500.50 AND sumRefund("ip", 1444) < 10500.50
+-> accept;