From 4058f0718eb61c31a487587791bcfb42c31b74f0 Mon Sep 17 00:00:00 2001
From: Piotr Przybylak <piotr.przybylak@gmail.com>
Date: Tue, 5 Mar 2024 21:07:30 +0100
Subject: [PATCH] Fix reporting for `Mockito.only` when wanted call is first
 (#3287)

Fixes #3286
---
 .../mockito/internal/verification/Only.java   |  4 +-
 .../internal/verification/OnlyTest.java       | 48 +++++++++++++++----
 2 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/src/main/java/org/mockito/internal/verification/Only.java b/src/main/java/org/mockito/internal/verification/Only.java
index aea8981d10..d5532b7088 100644
--- a/src/main/java/org/mockito/internal/verification/Only.java
+++ b/src/main/java/org/mockito/internal/verification/Only.java
@@ -25,6 +25,9 @@ public void verify(VerificationData data) {
         MatchableInvocation target = data.getTarget();
         List<Invocation> invocations = data.getAllInvocations();
         List<Invocation> chunk = findInvocations(invocations, target);
+        if (!chunk.isEmpty()) {
+            markVerified(chunk.get(0), target);
+        }
         if (invocations.size() != 1 && !chunk.isEmpty()) {
             Invocation unverified = findFirstUnverified(invocations);
             throw noMoreInteractionsWanted(unverified, (List) invocations);
@@ -32,7 +35,6 @@ public void verify(VerificationData data) {
         if (invocations.size() != 1 || chunk.isEmpty()) {
             throw wantedButNotInvoked(target);
         }
-        markVerified(chunk.get(0), target);
     }
 
     @Override
diff --git a/src/test/java/org/mockito/internal/verification/OnlyTest.java b/src/test/java/org/mockito/internal/verification/OnlyTest.java
index c269e3bff5..006173eac6 100644
--- a/src/test/java/org/mockito/internal/verification/OnlyTest.java
+++ b/src/test/java/org/mockito/internal/verification/OnlyTest.java
@@ -10,28 +10,34 @@
 import java.util.List;
 
 import org.junit.Test;
-import org.mockito.exceptions.base.MockitoAssertionError;
+import org.mockito.Mock;
+import org.mockito.exceptions.verification.NoInteractionsWanted;
+import org.mockito.exceptions.verification.WantedButNotInvoked;
 import org.mockito.internal.invocation.InvocationBuilder;
 import org.mockito.internal.invocation.InvocationMatcher;
 import org.mockito.internal.verification.api.VerificationData;
 import org.mockito.invocation.Invocation;
 import org.mockito.invocation.MatchableInvocation;
+import org.mockitousage.IMethods;
+import org.mockitoutil.TestBase;
 
-public class OnlyTest {
+public class OnlyTest extends TestBase {
+
+    @Mock IMethods mock;
 
     Only only = new Only();
 
-    public class VerificationDataStub implements VerificationData {
-        private final Invocation invocation;
+    public static class VerificationDataStub implements VerificationData {
+        private final Invocation[] invocations;
         private final InvocationMatcher wanted;
 
-        public VerificationDataStub(InvocationMatcher wanted, Invocation invocation) {
-            this.invocation = invocation;
+        public VerificationDataStub(InvocationMatcher wanted, Invocation... invocations) {
+            this.invocations = invocations;
             this.wanted = wanted;
         }
 
         public List<Invocation> getAllInvocations() {
-            return Arrays.asList(invocation);
+            return Arrays.asList(invocations);
         }
 
         @Override
@@ -45,7 +51,7 @@ public InvocationMatcher getWanted() {
     }
 
     @Test
-    public void shouldMarkAsVerified() {
+    public void shouldMarkAsVerifiedWhenAssertionSucceded() {
         // given
         Invocation invocation = new InvocationBuilder().toInvocation();
         assertFalse(invocation.isVerified());
@@ -58,7 +64,7 @@ public void shouldMarkAsVerified() {
     }
 
     @Test
-    public void shouldNotMarkAsVerifiedWhenAssertionFailed() {
+    public void shouldNotMarkAsVerifiedWhenWantedButNotInvoked() {
         // given
         Invocation invocation = new InvocationBuilder().toInvocation();
         assertFalse(invocation.isVerified());
@@ -69,10 +75,32 @@ public void shouldNotMarkAsVerifiedWhenAssertionFailed() {
                     new VerificationDataStub(
                             new InvocationBuilder().toInvocationMatcher(), invocation));
             fail();
-        } catch (MockitoAssertionError e) {
+        } catch (WantedButNotInvoked e) {
         }
 
         // then
         assertFalse(invocation.isVerified());
     }
+
+    @Test
+    public void shouldMarkWantedOnlyAsVerified() {
+
+        // given
+        InvocationBuilder invocationBuilder = new InvocationBuilder();
+        Invocation wanted = invocationBuilder.mock(mock).simpleMethod().toInvocation();
+        Invocation different = new InvocationBuilder().mock(mock).differentMethod().toInvocation();
+
+        // when
+        try {
+            only.verify(
+                    new VerificationDataStub(
+                            invocationBuilder.toInvocationMatcher(), wanted, different));
+            fail();
+        } catch (NoInteractionsWanted e) {
+        }
+
+        // then
+        assertTrue(wanted.isVerified());
+        assertFalse(different.isVerified());
+    }
 }