From cd70006bb6bccea7a9ce529cef35c8e46de65940 Mon Sep 17 00:00:00 2001
From: Haroon Khan <haroonkhan79@gmail.com>
Date: Thu, 17 Feb 2022 08:29:23 +0000
Subject: [PATCH] [JAVA-9986] Fix ThreadStamped unit test

---
 .../StampedAccount.java                       | 32 +++++++++++++------
 .../ThreadStampedAccountUnitTest.java         | 29 +++++++++++++----
 2 files changed, 45 insertions(+), 16 deletions(-)

diff --git a/core-java-modules/core-java-concurrency-advanced-3/src/main/java/com/baeldung/atomicstampedreference/StampedAccount.java b/core-java-modules/core-java-concurrency-advanced-3/src/main/java/com/baeldung/atomicstampedreference/StampedAccount.java
index 415b24738a..69aed0b979 100644
--- a/core-java-modules/core-java-concurrency-advanced-3/src/main/java/com/baeldung/atomicstampedreference/StampedAccount.java
+++ b/core-java-modules/core-java-concurrency-advanced-3/src/main/java/com/baeldung/atomicstampedreference/StampedAccount.java
@@ -1,25 +1,22 @@
 package com.baeldung.atomicstampedreference;
 
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicStampedReference;
 
 public class StampedAccount {
 
-    private AtomicInteger stamp = new AtomicInteger(0);
-    private AtomicStampedReference<Integer> account = new AtomicStampedReference<>(0, 0);
-
-    public int getBalance() {
-        return account.getReference();
-    }
-
-    public int getStamp() {
-        return account.getStamp();
-    }
+    private final AtomicInteger stamp = new AtomicInteger(0);
+    private final AtomicStampedReference<Integer> account = new AtomicStampedReference<>(0, 0);
 
     public boolean deposit(int funds) {
         int[] stamps = new int[1];
         int current = this.account.get(stamps);
         int newStamp = this.stamp.incrementAndGet();
+
+        // Thread is paused here to allow other threads to update the stamp and amount (for testing only)
+        sleep();
+
         return this.account.compareAndSet(current, current + funds, stamps[0], newStamp);
     }
 
@@ -29,4 +26,19 @@ public class StampedAccount {
         int newStamp = this.stamp.incrementAndGet();
         return this.account.compareAndSet(current, current - funds, stamps[0], newStamp);
     }
+
+    public int getBalance() {
+        return account.getReference();
+    }
+
+    public int getStamp() {
+        return account.getStamp();
+    }
+
+    private static void sleep() {
+        try {
+            TimeUnit.SECONDS.sleep(1);
+        } catch (InterruptedException ignored) {
+        }
+    }
 }
diff --git a/core-java-modules/core-java-concurrency-advanced-3/src/test/java/com/baeldung/atomicstampedreference/ThreadStampedAccountUnitTest.java b/core-java-modules/core-java-concurrency-advanced-3/src/test/java/com/baeldung/atomicstampedreference/ThreadStampedAccountUnitTest.java
index ce83355073..2840c25cf6 100644
--- a/core-java-modules/core-java-concurrency-advanced-3/src/test/java/com/baeldung/atomicstampedreference/ThreadStampedAccountUnitTest.java
+++ b/core-java-modules/core-java-concurrency-advanced-3/src/test/java/com/baeldung/atomicstampedreference/ThreadStampedAccountUnitTest.java
@@ -1,21 +1,38 @@
 package com.baeldung.atomicstampedreference;
 
-import org.junit.Assert;
 import org.junit.Test;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
 public class ThreadStampedAccountUnitTest {
 
     @Test
     public void givenMultiThread_whenStampedAccount_thenSetBalance() throws InterruptedException {
         StampedAccount account = new StampedAccount();
+
         Thread t = new Thread(() -> {
-            while (!account.withdrawal(100))
+            while (!account.deposit(100)) {
                 Thread.yield();
+            }
         });
         t.start();
-        Assert.assertTrue(account.deposit(100));
-        t.join(1_000);
-        Assert.assertFalse(t.isAlive());
-        Assert.assertSame(0, account.getBalance());
+
+        Thread t2 = new Thread(() -> {
+            while (!account.withdrawal(100)) {
+                Thread.yield();
+            }
+        });
+        t2.start();
+
+        t.join(10_000);
+        t2.join(10_000);
+
+        assertFalse(t.isAlive());
+        assertFalse(t2.isAlive());
+
+        assertEquals(0, account.getBalance());
+        assertTrue(account.getStamp() > 0);
     }
 }