From e5dbbaccd6008a44cdc3be22b6d705eab7f40d6c Mon Sep 17 00:00:00 2001
From: Pawel Kruszewski OpenSlava 2023 <pawel.kruszewski@ableneo.com>
Date: Thu, 19 Oct 2023 02:09:19 -0700
Subject: [PATCH] 4

---
 .../codefix/BalloonChallenge.java             | 24 +++++++++++--------
 .../codefix/BalloonChallengeTest.java         |  2 +-
 2 files changed, 15 insertions(+), 11 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..62cabae 100644
--- a/src/main/java/sk/ableneo/openslava2023/codefix/BalloonChallenge.java
+++ b/src/main/java/sk/ableneo/openslava2023/codefix/BalloonChallenge.java
@@ -5,30 +5,34 @@
 import sk.ableneo.openslava2023.donotfix.Product;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.LinkedList;
+import java.util.List;
 
 // TODO: fix the class
 public class BalloonChallenge {
 
-    public LinkedList findOrdersForProduct(Product p, boolean debug) {
-        ArrayList l = new ArrayList();
-        ArrayList list = getAllOrders();
+    public List<Order> findOrdersForProduct(Product p, boolean debug) {
+        List<Order> l = new ArrayList();
+        List<Order> list = getAllOrders();
+
         for (int i = 0; i < list.size(); i++) {
-            Order order = (Order) list.get(i);
+            Order order = list.get(i);
             boolean found = false;
-            if (order.getProducts().size() > 0) {
-                for (int j = 0; j <= order.getProducts().size(); j++) {
+            List<Product> products = order.getProducts();
+            if (products.size() > 0) {
+                for (int j = 0; j < order.getProducts().size(); j++) {
                     Product p2 = (Product) order.getProducts().get(j);
-                    found = (p2 == p);
+                    found = p2.equals(p);
                 }
                 if (found && order != null)
                     l.add(order);
             }
         }
-        return new LinkedList(l);
+        return l;
     }
 
-    public ArrayList getAllOrders() {
-        return new ArrayList<Object>(OrderRepository.getAllOrders());
+    public List<Order> getAllOrders() {
+        return new ArrayList<Order>(OrderRepository.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"))))