1313
1414package com .baidu .bifromq .basekv .store .range ;
1515
16+ import static org .mockito .ArgumentMatchers .eq ;
1617import static org .mockito .Mockito .mock ;
1718import static org .mockito .Mockito .times ;
1819import static org .mockito .Mockito .verify ;
@@ -59,7 +60,8 @@ public void setUp() {
5960
6061 @ Test
6162 public void awaitDone () {
62- when (range .toReseter (snapshot )).thenReturn (reseter );
63+ IKVReseter reseter = mock (IKVReseter .class );
64+ when (range .toReseter (eq (snapshot ))).thenReturn (reseter );
6365 KVRangeRestorer restorer = new KVRangeRestorer (snapshot , range , messenger , metricManager , executor , 10 );
6466
6567 assertTrue (restorer .awaitDone ().isDone ());
@@ -128,7 +130,7 @@ public void restoreFromWithError() {
128130 @ Test
129131 public void restoreFromTimeout () {
130132 IKVReseter reseter = mock (IKVReseter .class );
131- when (range .toReseter (snapshot )).thenReturn (reseter );
133+ when (range .toReseter (eq ( snapshot ) )).thenReturn (reseter );
132134
133135 KVRangeRestorer restorer = new KVRangeRestorer (snapshot , range , messenger , metricManager , executor , 1 );
134136 CompletableFuture <Void > restoreFuture = restorer .restoreFrom ("leader" , snapshot );
@@ -143,7 +145,7 @@ public void restoreFromTimeout() {
143145 @ Test
144146 public void cancelPreviousSession () {
145147 IKVReseter reseter = mock (IKVReseter .class );
146- when (range .toReseter (snapshot )).thenReturn (reseter );
148+ when (range .toReseter (eq ( snapshot ) )).thenReturn (reseter );
147149
148150 KVRangeRestorer restorer = new KVRangeRestorer (snapshot , range , messenger , metricManager , executor , 10 );
149151
@@ -152,6 +154,7 @@ public void cancelPreviousSession() {
152154
153155 // Start the second restore session, which should cancel the first
154156 KVRangeSnapshot newSnapshot = KVRangeSnapshot .newBuilder ().setId (snapshot .getId ()).setVer (1 ).build ();
157+ when (range .toReseter (eq (newSnapshot ))).thenReturn (reseter );
155158 CompletableFuture <Void > secondRestore = restorer .restoreFrom ("leader" , newSnapshot );
156159
157160 verify (reseter , times (1 )).abort ();
0 commit comments