From 777354428b62c413cdb8aa2f3c8f22737d5334b1 Mon Sep 17 00:00:00 2001
From: Pawel Kruszewski OpenSlava 2023 <pawel.kruszewski@ableneo.com>
Date: Thu, 19 Oct 2023 04:24:53 -0700
Subject: [PATCH] 10

---
 .../codefix/BalloonChallenge.java             | 26 +++++++++++--------
 .../codefix/BalloonChallengeTest.java         |  2 +-
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/src/main/java/sk/ableneo/openslava2023/codefix/BalloonChallenge.java b/src/main/java/sk/ableneo/openslava2023/codefix/BalloonChallenge.java
index 9b3e0a3..0bc8589 100644
--- a/src/main/java/sk/ableneo/openslava2023/codefix/BalloonChallenge.java
+++ b/src/main/java/sk/ableneo/openslava2023/codefix/BalloonChallenge.java
@@ -11,21 +11,25 @@
 public class BalloonChallenge {
 
     public LinkedList findOrdersForProduct(Product p, boolean debug) {
-        ArrayList l = new ArrayList();
-        ArrayList list = getAllOrders();
-        for (int i = 0; i < list.size(); i++) {
-            Order order = (Order) list.get(i);
+        ArrayList checkedOrderList = new ArrayList();
+        ArrayList allOrderslist = getAllOrders();
+        for (Object o : allOrderslist) {
+            Order order = (Order) o;
             boolean found = false;
-            if (order.getProducts().size() > 0) {
-                for (int j = 0; j <= order.getProducts().size(); j++) {
-                    Product p2 = (Product) order.getProducts().get(j);
-                    found = (p2 == p);
+            if (!order.getProducts().isEmpty()) {
+                for (int j = 0; j < order.getProducts().size(); j++) {
+                    Product p2 = order.getProducts().get(j);
+                    found = (p2.equals(p));
                 }
-                if (found && order != null)
-                    l.add(order);
+                if (found)
+                    try {
+                        checkedOrderList.add(order);
+                    } catch (RuntimeException e) {
+                        throw new RuntimeException("Can`t add order " + e);
+                    }
             }
         }
-        return new LinkedList(l);
+        return new LinkedList(checkedOrderList);
     }
 
     public ArrayList getAllOrders() {
diff --git a/src/test/java/sk/ableneo/openslava2023/codefix/BalloonChallengeTest.java b/src/test/java/sk/ableneo/openslava2023/codefix/BalloonChallengeTest.java
index 25c2fbe..b38003f 100644
--- a/src/test/java/sk/ableneo/openslava2023/codefix/BalloonChallengeTest.java
+++ b/src/test/java/sk/ableneo/openslava2023/codefix/BalloonChallengeTest.java
@@ -24,7 +24,7 @@ class BalloonChallengeTest {
     @BeforeEach
     void setUp() {
         when(balloonChallenge.getAllOrders()).thenReturn(
-                new ArrayList<Object>(Arrays.asList(
+                new ArrayList<Order>(Arrays.asList(
                         new Order(new ArrayList<>(Arrays.asList(new Product("hot air balloon")))),
                         new Order(new ArrayList<>(Arrays.asList(new Product("gas balloon")))),
                         new Order(new ArrayList<>(Arrays.asList(new Product("gas balloon"))))