Improve solution of ATM task

Author: andrei-punko

src/main/java/by/andd3dfx/common/ATM.java

@@ -25,10 +25,31 @@ public ATM(Map<Integer, Integer> state) {
                 .sorted(Comparator.reverseOrder()).toList();
     }
 
+    /**
+     * Withdraw asked amount using banknotes of ATM
+     *
+     * @param amount sum asked to withdraw
+     * @return map with solution - pairs {banknote nominal->quantity}
+     */
     public Map<Integer, Integer> withdraw(int amount) {
+        // Try to make withdraw using banknote of highest nominal,
+        // in case of fail - try to start from next nominal
+        for (int startingBanknoteIndex = 0; startingBanknoteIndex < nominals.size(); startingBanknoteIndex++) {
+            var result = withdraw(amount, startingBanknoteIndex);
+
+            if (result != null) {
+                return result;
+            }
+        }
+
+        throw new IllegalStateException("Could not perform withdraw!");
+    }
+
+    private Map<Integer, Integer> withdraw(int amount, int startingBanknoteIndex) {
         var result = new HashMap<Integer, Integer>();
 
-        for (var nominal : nominals) {
+        for (var index = startingBanknoteIndex; index < nominals.size(); index++) {
+            var nominal = nominals.get(index);
             if (nominal > amount || state.get(nominal) == 0) {
                 continue;
             }
@@ -45,7 +66,7 @@ public Map<Integer, Integer> withdraw(int amount) {
         }
 
         if (amount > 0) {
-            throw new IllegalStateException("Could not perform withdraw!");
+            return null;
         }
 
         for (var nominal : result.keySet()) {

src/test/java/by/andd3dfx/common/ATMTest.java

@@ -16,8 +16,8 @@ public class ATMTest {
     public void setUp() throws Exception {
         atm = new ATM(Map.of(
                 500, 1,
-                100, 2,
-                50, 10
+                200, 3,
+                50, 5
         ));
     }
 
@@ -26,8 +26,23 @@ public void withdraw() {
         var result = atm.withdraw(450);
 
         assertThat(result).isEqualTo(Map.of(
-                100, 2,
-                50, 5
+                200, 2,
+                50, 1
+        ));
+    }
+
+    @Test
+    public void withdrawWhenHighGradeBanknotePresentButShouldNotBeUsed() {
+        // Spend all 50 banknotes
+        for (int i = 0; i < 5; i++) {
+            atm.withdraw(50);
+        }
+
+        // Ask 600, ATM has only 500 & 200 banknotes at this moment
+        var result = atm.withdraw(600);
+
+        assertThat(result).isEqualTo(Map.of(
+                200, 3
         ));
     }
 
@@ -42,14 +57,13 @@ public void withdrawForConsequentCalls() {
         var result = atm.withdraw(650);
         assertThat(result).isEqualTo(Map.of(
                 500, 1,
-                100, 1,
-                50, 1
+                50, 3
         ));
 
-        var result2 = atm.withdraw(250);
+        var result2 = atm.withdraw(650);
         assertThat(result2).isEqualTo(Map.of(
-                100, 1,
-                50, 3
+                200, 3,
+                50, 1
         ));
     }
 }