TestSubscriber - fix awaitTerminalEventAndUnsubscribeOnTimeout

src/main/java/rx/observers/TestSubscriber.java

@@ -281,7 +281,7 @@ public void awaitTerminalEvent(long timeout, TimeUnit unit) {
      * Blocks until this {@link Subscriber} receives a notification that the {@code Observable} is complete
      * (either an {@code onCompleted} or {@code onError} notification), or until a timeout expires; if the
      * Subscriber is interrupted before either of these events take place, this method unsubscribes the
-     * Subscriber from the Observable).
+     * Subscriber from the Observable). If timeout expires then the Subscriber is unsubscribed from the Observable.
      *
      * @param timeout
      *          the duration of the timeout
@@ -290,8 +290,12 @@ public void awaitTerminalEvent(long timeout, TimeUnit unit) {
      */
     public void awaitTerminalEventAndUnsubscribeOnTimeout(long timeout, TimeUnit unit) {
         try {
-            awaitTerminalEvent(timeout, unit);
-        } catch (RuntimeException e) {
+            boolean result = latch.await(timeout, unit);
+            if (!result) {
+                // timeout occurred
+                unsubscribe();
+            }
+        } catch (InterruptedException e) {
             unsubscribe();
         }
     }

src/test/java/rx/observers/TestSubscriberTest.java

@@ -16,19 +16,24 @@
 package rx.observers;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 
 import java.util.Arrays;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
+import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 import org.mockito.InOrder;
 
 import rx.Observable;
 import rx.Observer;
+import rx.functions.Action0;
 import rx.subjects.PublishSubject;
 
 public class TestSubscriberTest {
@@ -124,8 +129,35 @@ public void testWrappingMockWhenUnsubscribeInvolved() {
     @Test
     public void testAssertError() {
         RuntimeException e = new RuntimeException("Oops");
-        TestSubscriber subscriber = new TestSubscriber();
+        TestSubscriber<Object> subscriber = new TestSubscriber<Object>();
         Observable.error(e).subscribe(subscriber);
         subscriber.assertError(e);
     }
+
+    @Test
+    public void testAwaitTerminalEventWithDuration() {
+        TestSubscriber<Object> ts = new TestSubscriber<Object>();
+        Observable.just(1).subscribe(ts);
+        ts.awaitTerminalEvent(1, TimeUnit.SECONDS);
+        ts.assertTerminalEvent();
+    }
+
+    @Test
+    public void testAwaitTerminalEventWithDurationAndUnsubscribeOnTimeout() {
+        TestSubscriber<Object> ts = new TestSubscriber<Object>();
+        final AtomicBoolean unsub = new AtomicBoolean(false);
+        Observable.just(1)
+        //
+                .doOnUnsubscribe(new Action0() {
+                    @Override
+                    public void call() {
+                        unsub.set(true);
+                    }
+                })
+                //
+                .delay(1000, TimeUnit.MILLISECONDS).subscribe(ts);
+        ts.awaitTerminalEventAndUnsubscribeOnTimeout(100, TimeUnit.MILLISECONDS);
+        assertTrue(unsub.get());
+    }
+
 }