From c605a711a6b94005c7f61facab66d4f6da57ff06 Mon Sep 17 00:00:00 2001
From: Mark Struchkov <upagge@icloud.com>
Date: Wed, 21 Aug 2019 20:55:54 +0300
Subject: [PATCH] Optional

---
 .../sadtech/autoresponder/AutoResponder.java  | 59 +++++++------------
 .../compare/UnitPriorityComparator.java       |  2 +-
 .../sadtech/autoresponder/entity/Unit.java    |  7 ++-
 .../autoresponder/entity/UnitPointer.java     | 12 ++--
 .../repository/UnitPointerRepository.java     | 15 ++---
 .../repository/UnitPointerRepositoryMap.java  | 17 +++---
 .../service/UnitPointerService.java           | 16 +++--
 .../service/UnitPointerServiceImpl.java       | 27 +++++----
 .../autoresponder/AutoResponderTest.java      | 45 +++++++-------
 .../sadtech/autoresponder/test/TestUnit.java  |  4 +-
 10 files changed, 101 insertions(+), 103 deletions(-)

diff --git a/src/main/java/org/sadtech/autoresponder/AutoResponder.java b/src/main/java/org/sadtech/autoresponder/AutoResponder.java
index 3b87314..e94b3c3 100644
--- a/src/main/java/org/sadtech/autoresponder/AutoResponder.java
+++ b/src/main/java/org/sadtech/autoresponder/AutoResponder.java
@@ -7,7 +7,6 @@ import lombok.extern.slf4j.Slf4j;
 import org.sadtech.autoresponder.compare.UnitPriorityComparator;
 import org.sadtech.autoresponder.entity.Unit;
 import org.sadtech.autoresponder.entity.UnitPointer;
-import org.sadtech.autoresponder.exception.NotFoundUnitException;
 import org.sadtech.autoresponder.service.UnitPointerService;
 import org.sadtech.autoresponder.util.Description;
 import org.sadtech.autoresponder.util.Parser;
@@ -25,7 +24,7 @@ import java.util.Set;
 public class AutoResponder<U extends Unit> {
 
     @Description("Компоратор для сортировки Unit-ов")
-    private final UnitPriorityComparator<U> unitPriorityComparator = new UnitPriorityComparator<>();
+    private static final UnitPriorityComparator UNIT_PRIORITY_COMPARATOR = new UnitPriorityComparator();
 
     @Description("Начальные юниты, первый запрос приходит на них")
     private final Set<U> startUnits;
@@ -36,9 +35,9 @@ public class AutoResponder<U extends Unit> {
     private U defaultUnit;
 
     @Getter
-    private final UnitPointerService<U> unitPointerService;
+    private final UnitPointerService unitPointerService;
 
-    public AutoResponder(UnitPointerService<U> unitPointerService, Set<U> startUnits) {
+    public AutoResponder(UnitPointerService unitPointerService, Set<U> startUnits) {
         this.startUnits = startUnits;
         this.unitPointerService = unitPointerService;
     }
@@ -50,20 +49,21 @@ public class AutoResponder<U extends Unit> {
      * @param message  Запрос пользователя - текстовое сообщение
      * @return {@link Unit}, который отвечает за данные для обработки данного запроса
      */
-    public U answer(@NonNull Integer personId, String message) {
-        UnitPointer<U> unitPointer = checkAndAddPerson(personId);
-        U unit;
-        try {
-            if (unitPointer.getUnit() == null || unitPointer.getUnit().getNextUnits() == null || unitPointer.getUnit().getNextUnits().isEmpty()) {
-                unit = nextUnit(startUnits, message);
-            } else {
-                unit = (U) nextUnit(unitPointer.getUnit().getNextUnits(), message);
-            }
+    public Optional<U> answer(@NonNull Integer personId, String message) {
+        Optional<UnitPointer> unitPointer = unitPointerService.getByEntityId(personId);
+        Optional<U> unitOpt = nextUnit(
+                !unitPointer.isPresent() || newScenario(unitPointer.get()) ? startUnits : unitPointer.get().getUnit().getNextUnits(),
+                message
+        );
+        unitOpt.ifPresent(unit -> {
+            unitPointer.orElseGet(() -> unitPointerService.create(personId));
             unitPointerService.edit(personId, unit);
-        } catch (NotFoundUnitException e) {
-            unit = defaultUnit;
-        }
-        return Optional.ofNullable(unit).orElseThrow(NotFoundUnitException::new);
+        });
+        return unitOpt.isPresent() ? unitOpt : Optional.ofNullable(defaultUnit);
+    }
+
+    private boolean newScenario(UnitPointer unitPointer) {
+        return unitPointer.getUnit() == null || unitPointer.getUnit().getNextUnits() == null || unitPointer.getUnit().getNextUnits().isEmpty();
     }
 
     /**
@@ -73,7 +73,7 @@ public class AutoResponder<U extends Unit> {
      * @param message   Запрос пользователя - текстовое сообщение
      * @return Юнит, который нуждается в обработке в соответствии с запросом пользователя
      */
-    private U nextUnit(@NonNull Set<U> nextUnits, String message) {
+    private Optional<U> nextUnit(@NonNull Set<U> nextUnits, String message) {
         Set<U> searchUnit = new HashSet<>();
 
         nextUnits.stream()
@@ -97,32 +97,15 @@ public class AutoResponder<U extends Unit> {
                     .forEach(searchUnit::add);
         }
 
-        return searchUnit.stream().max(unitPriorityComparator).orElseThrow(NotFoundUnitException::new);
-    }
-
-    /**
-     * Проверяет наличие {@link UnitPointer} у пользовтеля и создает его, если не находит
-     *
-     * @param personId Идентификатор пользователя
-     * @return {@link UnitPointer}, который сохраняет {@link Unit}, который был обработан последним
-     */
-    private UnitPointer<U> checkAndAddPerson(@NonNull Integer personId) {
-        UnitPointer<U> unitPointer;
-        if (unitPointerService.check(personId)) {
-            unitPointer = unitPointerService.getByEntityId(personId);
-        } else {
-            unitPointer = new UnitPointer<>(personId);
-            unitPointerService.add(unitPointer);
-        }
-        return unitPointer;
+        return searchUnit.stream().max(UNIT_PRIORITY_COMPARATOR);
     }
 
 
-    private boolean patternReg(@NonNull Unit unit, String message) {
+    private boolean patternReg(@NonNull U unit, String message) {
         return message.matches(unit.getPattern().pattern());
     }
 
-    private Double percentageMatch(Unit unit, Set<String> words) {
+    private Double percentageMatch(U unit, Set<String> words) {
         if (unit.getKeyWords() != null && !unit.getKeyWords().isEmpty()) {
             Set<String> temp = new HashSet<>(unit.getKeyWords());
             temp.retainAll(words);
diff --git a/src/main/java/org/sadtech/autoresponder/compare/UnitPriorityComparator.java b/src/main/java/org/sadtech/autoresponder/compare/UnitPriorityComparator.java
index 8fc1421..7b13c47 100644
--- a/src/main/java/org/sadtech/autoresponder/compare/UnitPriorityComparator.java
+++ b/src/main/java/org/sadtech/autoresponder/compare/UnitPriorityComparator.java
@@ -9,7 +9,7 @@ import java.util.Comparator;
  *
  * @author upagge [07/07/2019]
  */
-public class UnitPriorityComparator<U extends Unit> implements Comparator<U> {
+public class UnitPriorityComparator implements Comparator<Unit> {
 
     @Override
     public int compare(Unit o1, Unit o2) {
diff --git a/src/main/java/org/sadtech/autoresponder/entity/Unit.java b/src/main/java/org/sadtech/autoresponder/entity/Unit.java
index c00ed4f..cf8900e 100644
--- a/src/main/java/org/sadtech/autoresponder/entity/Unit.java
+++ b/src/main/java/org/sadtech/autoresponder/entity/Unit.java
@@ -39,7 +39,12 @@ public abstract class Unit<U extends Unit> {
     @Description("Множество следующих Unit в сценарии")
     protected Set<U> nextUnits;
 
-    protected Unit(Set<String> keyWords, String phrase, Pattern pattern, Integer matchThreshold, Integer priority, Set<U> nextUnits) {
+    protected Unit(Set<String> keyWords,
+                   String phrase,
+                   Pattern pattern,
+                   Integer matchThreshold,
+                   Integer priority,
+                   Set<U> nextUnits) {
         this.keyWords = keyWords;
         this.phrase = phrase;
         this.pattern = pattern;
diff --git a/src/main/java/org/sadtech/autoresponder/entity/UnitPointer.java b/src/main/java/org/sadtech/autoresponder/entity/UnitPointer.java
index 05aec5f..a62dff5 100644
--- a/src/main/java/org/sadtech/autoresponder/entity/UnitPointer.java
+++ b/src/main/java/org/sadtech/autoresponder/entity/UnitPointer.java
@@ -2,6 +2,7 @@ package org.sadtech.autoresponder.entity;
 
 import lombok.AllArgsConstructor;
 import lombok.Data;
+import lombok.RequiredArgsConstructor;
 import org.sadtech.autoresponder.util.Description;
 
 /**
@@ -10,17 +11,14 @@ import org.sadtech.autoresponder.util.Description;
  * @author upagge [07/07/2019]
  */
 @Data
+@RequiredArgsConstructor
 @AllArgsConstructor
-public class UnitPointer<U extends Unit> {
+public class UnitPointer {
 
     @Description("Идентификатор пользователя")
-    private Integer entityId;
+    private final Integer entityId;
 
     @Description("Юнит, который был обработан")
-    private U unit;
-
-    public UnitPointer(Integer entityId) {
-        this.entityId = entityId;
-    }
+    private Unit unit;
 
 }
diff --git a/src/main/java/org/sadtech/autoresponder/repository/UnitPointerRepository.java b/src/main/java/org/sadtech/autoresponder/repository/UnitPointerRepository.java
index a889c52..1b65f50 100644
--- a/src/main/java/org/sadtech/autoresponder/repository/UnitPointerRepository.java
+++ b/src/main/java/org/sadtech/autoresponder/repository/UnitPointerRepository.java
@@ -1,29 +1,30 @@
 package org.sadtech.autoresponder.repository;
 
-import org.sadtech.autoresponder.entity.Unit;
+import lombok.NonNull;
 import org.sadtech.autoresponder.entity.UnitPointer;
 
 import java.util.Collection;
+import java.util.Optional;
 
 /**
  * Интегрфейс для работы с хранилищем сущности {@link UnitPointer}.
  *
  * @author upagge [07/07/2019]
  */
-public interface UnitPointerRepository<U extends Unit> {
+public interface UnitPointerRepository {
 
-    void add(UnitPointer<U> unitPointer);
+    UnitPointer add(@NonNull UnitPointer unitPointer);
 
-    void edit(UnitPointer<U> unitPointer);
+    void edit(@NonNull UnitPointer unitPointer);
 
-    void remove(Integer entityId);
+    void remove(@NonNull Integer entityId);
 
-    void addAll(Collection<UnitPointer<U>> unitPointers);
+    void addAll(@NonNull Collection<UnitPointer> unitPointers);
 
     /**
      * @param entityId Идентификатор пользователя
      * @return Объект с последним обработанным {@link org.sadtech.autoresponder.entity.Unit} для пользователя
      */
-    UnitPointer<U> findByEntityId(Integer entityId);
+    Optional<UnitPointer> findByEntityId(@NonNull Integer entityId);
 
 }
diff --git a/src/main/java/org/sadtech/autoresponder/repository/UnitPointerRepositoryMap.java b/src/main/java/org/sadtech/autoresponder/repository/UnitPointerRepositoryMap.java
index 1c22747..af25cb1 100644
--- a/src/main/java/org/sadtech/autoresponder/repository/UnitPointerRepositoryMap.java
+++ b/src/main/java/org/sadtech/autoresponder/repository/UnitPointerRepositoryMap.java
@@ -1,29 +1,30 @@
 package org.sadtech.autoresponder.repository;
 
 import lombok.NonNull;
-import org.sadtech.autoresponder.entity.Unit;
 import org.sadtech.autoresponder.entity.UnitPointer;
 
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 
 /**
  * Реализация хранилища для {@link UnitPointer} на основе Map.
  *
  * @author upagge [07/07/2019]
  */
-public class UnitPointerRepositoryMap<U extends Unit> implements UnitPointerRepository<U> {
+public class UnitPointerRepositoryMap implements UnitPointerRepository {
 
-    private Map<Integer, UnitPointer<U>> unitPointerMap = new HashMap<>();
+    private Map<Integer, UnitPointer> unitPointerMap = new HashMap<>();
 
     @Override
-    public void add(@NonNull UnitPointer<U> unitPointer) {
+    public UnitPointer add(@NonNull UnitPointer unitPointer) {
         unitPointerMap.put(unitPointer.getEntityId(), unitPointer);
+        return unitPointer;
     }
 
     @Override
-    public void edit(@NonNull UnitPointer<U> unitPointer) {
+    public void edit(@NonNull UnitPointer unitPointer) {
         unitPointerMap.get(unitPointer.getEntityId()).setUnit(unitPointer.getUnit());
     }
 
@@ -33,12 +34,12 @@ public class UnitPointerRepositoryMap<U extends Unit> implements UnitPointerRepo
     }
 
     @Override
-    public void addAll(@NonNull Collection<UnitPointer<U>> unitPointers) {
+    public void addAll(@NonNull Collection<UnitPointer> unitPointers) {
         unitPointers.parallelStream().forEach(unitPointer -> unitPointerMap.put(unitPointer.getEntityId(), unitPointer));
     }
 
     @Override
-    public UnitPointer<U> findByEntityId(@NonNull Integer entityId) {
-        return unitPointerMap.get(entityId);
+    public Optional<UnitPointer> findByEntityId(@NonNull Integer entityId) {
+        return Optional.ofNullable(unitPointerMap.get(entityId));
     }
 }
diff --git a/src/main/java/org/sadtech/autoresponder/service/UnitPointerService.java b/src/main/java/org/sadtech/autoresponder/service/UnitPointerService.java
index f134f54..06a55ca 100644
--- a/src/main/java/org/sadtech/autoresponder/service/UnitPointerService.java
+++ b/src/main/java/org/sadtech/autoresponder/service/UnitPointerService.java
@@ -1,16 +1,19 @@
 package org.sadtech.autoresponder.service;
 
+import lombok.NonNull;
 import org.sadtech.autoresponder.entity.Unit;
 import org.sadtech.autoresponder.entity.UnitPointer;
 
+import java.util.Optional;
+
 /**
  * Сервис для взаимодействия с сущностью {@link UnitPointer}.
  *
  * @author upagge [07/07/2019]
  */
-public interface UnitPointerService<U extends Unit> {
+public interface UnitPointerService {
 
-    void add(UnitPointer<U> unitPointer);
+    void add(@NonNull UnitPointer unitPointer);
 
     /**
      * Проверка наличия {@link UnitPointer} для пользователя
@@ -18,10 +21,13 @@ public interface UnitPointerService<U extends Unit> {
      * @param entityId Идентификатор пользователя
      * @return true - если найдено
      */
-    boolean check(Integer entityId);
+    boolean check(@NonNull Integer entityId);
 
-    UnitPointer<U> getByEntityId(Integer entityId);
+    Optional<UnitPointer> getByEntityId(@NonNull Integer entityId);
+
+    UnitPointer create(@NonNull Integer entryId);
+
+    void edit(@NonNull Integer personId, Unit unit);
 
-    void edit(Integer personId, U unit);
 
 }
diff --git a/src/main/java/org/sadtech/autoresponder/service/UnitPointerServiceImpl.java b/src/main/java/org/sadtech/autoresponder/service/UnitPointerServiceImpl.java
index 49e7150..c608627 100644
--- a/src/main/java/org/sadtech/autoresponder/service/UnitPointerServiceImpl.java
+++ b/src/main/java/org/sadtech/autoresponder/service/UnitPointerServiceImpl.java
@@ -1,40 +1,45 @@
 package org.sadtech.autoresponder.service;
 
 import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.sadtech.autoresponder.entity.Unit;
 import org.sadtech.autoresponder.entity.UnitPointer;
 import org.sadtech.autoresponder.repository.UnitPointerRepository;
 
+import java.util.Optional;
+
 @Slf4j
-public class UnitPointerServiceImpl<U extends Unit> implements UnitPointerService<U> {
+@RequiredArgsConstructor
+public class UnitPointerServiceImpl implements UnitPointerService {
 
-    private final UnitPointerRepository<U> unitPointerRepository;
-
-    public UnitPointerServiceImpl(UnitPointerRepository<U> unitPointerRepository) {
-        this.unitPointerRepository = unitPointerRepository;
-    }
+    private final UnitPointerRepository unitPointerRepository;
 
     @Override
-    public UnitPointer<U> getByEntityId(@NonNull Integer entityId) {
+    public Optional<UnitPointer> getByEntityId(@NonNull Integer entityId) {
         return unitPointerRepository.findByEntityId(entityId);
     }
 
     @Override
-    public void edit(@NonNull Integer personId, U unit) {
+    public UnitPointer create(@NonNull Integer entryId) {
+        return unitPointerRepository.add(new UnitPointer(entryId));
+    }
+
+    @Override
+    public void edit(@NonNull Integer personId, Unit unit) {
         if (check(personId)) {
-            unitPointerRepository.edit(new UnitPointer<>(personId, unit));
+            unitPointerRepository.edit(new UnitPointer(personId, unit));
         }
     }
 
     @Override
-    public void add(@NonNull UnitPointer<U> unitPointer) {
+    public void add(@NonNull UnitPointer unitPointer) {
         unitPointerRepository.add(unitPointer);
         log.info("Пользователь отправлен в репозиторий");
     }
 
     @Override
     public boolean check(@NonNull Integer entityId) {
-        return unitPointerRepository.findByEntityId(entityId) != null;
+        return unitPointerRepository.findByEntityId(entityId).isPresent();
     }
 }
diff --git a/src/test/java/org/sadtech/autoresponder/AutoResponderTest.java b/src/test/java/org/sadtech/autoresponder/AutoResponderTest.java
index b2a37fb..672908c 100644
--- a/src/test/java/org/sadtech/autoresponder/AutoResponderTest.java
+++ b/src/test/java/org/sadtech/autoresponder/AutoResponderTest.java
@@ -3,13 +3,12 @@ package org.sadtech.autoresponder;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
-import org.sadtech.autoresponder.entity.Unit;
-import org.sadtech.autoresponder.exception.NotFoundUnitException;
 import org.sadtech.autoresponder.repository.UnitPointerRepositoryMap;
 import org.sadtech.autoresponder.service.UnitPointerServiceImpl;
 import org.sadtech.autoresponder.test.TestUnit;
 
 import java.util.HashSet;
+import java.util.Optional;
 import java.util.Set;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
@@ -17,7 +16,7 @@ import java.util.stream.Stream;
 
 public class AutoResponderTest {
 
-    private AutoResponder autoresponder;
+    private AutoResponder<TestUnit> autoresponder;
 
     @Before
     public void setUp() {
@@ -60,20 +59,20 @@ public class AutoResponderTest {
                 .matchThreshold(100)
                 .build();
 
-        Set<Unit> testUnits = new HashSet<>();
+        Set<TestUnit> testUnits = new HashSet<>();
         testUnits.add(hello);
         testUnits.add(regExp);
         testUnits.add(unreal);
 
         UnitPointerServiceImpl unitPointerService = new UnitPointerServiceImpl(new UnitPointerRepositoryMap());
-        autoresponder = new AutoResponder(unitPointerService, testUnits);
+        autoresponder = new AutoResponder<>(unitPointerService, testUnits);
     }
 
     @Test
     public void simpleAnswer() {
-        String message = ((TestUnit) autoresponder.answer(1, "привет это тестирвоание функциональности")).getMessage();
+        String message = autoresponder.answer(1, "привет это тестирвоание функциональности").get().getMessage();
         Assert.assertEquals("тест", message);
-        String message1 = ((TestUnit) autoresponder.answer(2, "привет, еще одно тестирование")).getMessage();
+        String message1 = autoresponder.answer(2, "привет, еще одно тестирование").get().getMessage();
         Assert.assertEquals("тест", message1);
     }
 
@@ -82,17 +81,17 @@ public class AutoResponderTest {
         TestUnit defaultUnit = TestUnit.builder().message("не знаю").build();
         autoresponder.setDefaultUnit(defaultUnit);
 
-        Assert.assertEquals("не знаю", ((TestUnit) autoresponder.answer(2, "пока")).getMessage());
+        Assert.assertEquals("не знаю", autoresponder.answer(2, "пока").get().getMessage());
     }
 
-    @Test(expected = NotFoundUnitException.class)
+    @Test
     public void notDefaultAnswer() {
-        autoresponder.answer(2, "пока");
+        Assert.assertEquals(Optional.empty(), autoresponder.answer(2, "пока"));
     }
 
     @Test
     public void regExpAnswer() {
-        String message = ((TestUnit) autoresponder.answer(1, "79101234567")).getMessage();
+        String message = autoresponder.answer(1, "79101234567").get().getMessage();
         Assert.assertEquals("регулярка", message);
     }
 
@@ -100,9 +99,9 @@ public class AutoResponderTest {
     public void priorityAnswer() {
         autoresponder.answer(1, "привет");
         autoresponder.answer(2, "привет");
-        String message = ((TestUnit) autoresponder.answer(1, "как твои делишки")).getMessage();
+        String message = autoresponder.answer(1, "как твои делишки").get().getMessage();
         Assert.assertEquals("отлично", message);
-        String message1 = ((TestUnit) autoresponder.answer(2, "твои дела все еще хорошо?")).getMessage();
+        String message1 = autoresponder.answer(2, "твои дела все еще хорошо?").get().getMessage();
         Assert.assertEquals("хорошо", message1);
     }
 
@@ -110,14 +109,14 @@ public class AutoResponderTest {
     public void showRegExpVsKeyWords() {
         autoresponder.answer(1, "привет");
         autoresponder.answer(1, "дела");
-        Assert.assertEquals("регулярка", ((TestUnit) autoresponder.answer(1, "89101234567")).getMessage());
+        Assert.assertEquals("регулярка", autoresponder.answer(1, "89101234567").get().getMessage());
     }
 
-    @Test(expected = NotFoundUnitException.class)
+    @Test
     public void matchThreshold() {
         autoresponder.answer(1, "витамин я сьем, и арбуз получу");
         String message = "параметр себе задам: покушать витамин и арбуз, а вместе все это мультифрукт";
-        String answer = ((TestUnit) autoresponder.answer(1, message)).getMessage();
+        String answer = autoresponder.answer(1, message).get().getMessage();
         Assert.assertEquals("победа", answer);
     }
 
@@ -125,16 +124,16 @@ public class AutoResponderTest {
     public void generalAnswer() {
         TestUnit defaultUnit = TestUnit.builder().message("не знаю").build();
         autoresponder.setDefaultUnit(defaultUnit);
-        String answer = ((TestUnit) autoresponder.answer(1, "привет это тестирование функциональности")).getMessage();
+        String answer = autoresponder.answer(1, "привет это тестирование функциональности").get().getMessage();
         Assert.assertEquals("тест", answer);
-        Assert.assertEquals("хорошо", ((TestUnit) autoresponder.answer(1, "как твои дела")).getMessage());
-        String answer2 = ((TestUnit) autoresponder.answer(2, "привет это тестирование функциональности")).getMessage();
+        Assert.assertEquals("хорошо", autoresponder.answer(1, "как твои дела").get().getMessage());
+        String answer2 = autoresponder.answer(2, "привет это тестирование функциональности").get().getMessage();
         Assert.assertEquals("тест", answer2);
-        Assert.assertEquals("не знаю", ((TestUnit) autoresponder.answer(1, "нет")).getMessage());
+        Assert.assertEquals("не знаю", autoresponder.answer(1, "нет").get().getMessage());
         String message = "параметр себе задам: покушать витамин и арбуз, а вместе все это мультифрукт";
-        Assert.assertEquals("победа", ((TestUnit) autoresponder.answer(3, message)).getMessage());
-        Assert.assertEquals("регулярка", ((TestUnit) autoresponder.answer(1, "8912345678")).getMessage());
-        String answer3 = ((TestUnit) autoresponder.answer(1, "привет это тестирование функциональности")).getMessage();
+        Assert.assertEquals("победа", autoresponder.answer(3, message).get().getMessage());
+        Assert.assertEquals("регулярка", autoresponder.answer(1, "8912345678").get().getMessage());
+        String answer3 = autoresponder.answer(1, "привет это тестирование функциональности").get().getMessage();
         Assert.assertEquals("тест", answer3);
     }
 
diff --git a/src/test/java/org/sadtech/autoresponder/test/TestUnit.java b/src/test/java/org/sadtech/autoresponder/test/TestUnit.java
index 79caeae..ee14e9b 100644
--- a/src/test/java/org/sadtech/autoresponder/test/TestUnit.java
+++ b/src/test/java/org/sadtech/autoresponder/test/TestUnit.java
@@ -13,7 +13,7 @@ import java.util.regex.Pattern;
 @Getter
 @EqualsAndHashCode(callSuper = true)
 @ToString
-public class TestUnit extends Unit {
+public class TestUnit extends Unit<TestUnit> {
 
     private String message;
 
@@ -23,7 +23,7 @@ public class TestUnit extends Unit {
                     Pattern pattern,
                     Integer matchThreshold,
                     Integer priority,
-                    @Singular Set<Unit> nextUnits,
+                    @Singular Set<TestUnit> nextUnits,
                     String message) {
         super(keyWords, phrase, pattern, matchThreshold, priority, nextUnits);
         this.message = message;