From a0bc6c9b2ea04356d7b953fdfed1c76c09352473 Mon Sep 17 00:00:00 2001
From: Aliaksei Breilian <aliaksei_breilian@epam.com>
Date: Tue, 15 Feb 2022 12:00:07 +0300
Subject: [PATCH] Update CodeBuilder.java

---
 src/test/java/oops/CodeBuilder.java | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/src/test/java/oops/CodeBuilder.java b/src/test/java/oops/CodeBuilder.java
index 56d9c30..432b29f 100644
--- a/src/test/java/oops/CodeBuilder.java
+++ b/src/test/java/oops/CodeBuilder.java
@@ -53,6 +53,32 @@ private Map<String, List<Steps>> separateSteps(List<Steps> steps) {
         return allPairs;
     }
 
+    
+    private static Map<String, Collection<Steps>> separateSteps(List<Steps> steps) {
+        //import com.google.common.collect.LinkedListMultimap;
+        //import com.google.common.collect.Multimap;
+        Multimap<String, Steps> result = LinkedListMultimap.create();
+        AtomicReference<String> currentKey = new AtomicReference<>();
+
+        if (!steps.get(0).getClassName().contains("Reports.logCase")) {
+            currentKey.set("DefaultCaseId");
+        }
+
+        steps.forEach(s -> {
+            if (s.getClassName().contains("Reports.logCase")) {
+                currentKey.set(s.getCalssName()); // actualy we should get step parameter
+            } else {
+                result.put(currentKey.get(), s);
+            }
+        });
+        
+        if (!result.containsKey(currentKey.get())) {
+            throws new RuntimeException("unexpected report step at the end of test");
+        }
+
+        return result.asMap();
+    }
+
     private List<MethodSpec> buildMethods(Map<String, List<Steps>> allPairs, String className) {
         List<MethodSpec> methods = new ArrayList<>();
         allPairs.forEach((annotation, steps) -> {