13
13
14
14
package com .baidu .bifromq .basekv .store .range ;
15
15
16
+ import static org .mockito .ArgumentMatchers .eq ;
16
17
import static org .mockito .Mockito .mock ;
17
18
import static org .mockito .Mockito .times ;
18
19
import static org .mockito .Mockito .verify ;
@@ -59,7 +60,8 @@ public void setUp() {
59
60
60
61
@ Test
61
62
public void awaitDone () {
62
- when (range .toReseter (snapshot )).thenReturn (reseter );
63
+ IKVReseter reseter = mock (IKVReseter .class );
64
+ when (range .toReseter (eq (snapshot ))).thenReturn (reseter );
63
65
KVRangeRestorer restorer = new KVRangeRestorer (snapshot , range , messenger , metricManager , executor , 10 );
64
66
65
67
assertTrue (restorer .awaitDone ().isDone ());
@@ -128,7 +130,7 @@ public void restoreFromWithError() {
128
130
@ Test
129
131
public void restoreFromTimeout () {
130
132
IKVReseter reseter = mock (IKVReseter .class );
131
- when (range .toReseter (snapshot )).thenReturn (reseter );
133
+ when (range .toReseter (eq ( snapshot ) )).thenReturn (reseter );
132
134
133
135
KVRangeRestorer restorer = new KVRangeRestorer (snapshot , range , messenger , metricManager , executor , 1 );
134
136
CompletableFuture <Void > restoreFuture = restorer .restoreFrom ("leader" , snapshot );
@@ -143,7 +145,7 @@ public void restoreFromTimeout() {
143
145
@ Test
144
146
public void cancelPreviousSession () {
145
147
IKVReseter reseter = mock (IKVReseter .class );
146
- when (range .toReseter (snapshot )).thenReturn (reseter );
148
+ when (range .toReseter (eq ( snapshot ) )).thenReturn (reseter );
147
149
148
150
KVRangeRestorer restorer = new KVRangeRestorer (snapshot , range , messenger , metricManager , executor , 10 );
149
151
@@ -152,6 +154,7 @@ public void cancelPreviousSession() {
152
154
153
155
// Start the second restore session, which should cancel the first
154
156
KVRangeSnapshot newSnapshot = KVRangeSnapshot .newBuilder ().setId (snapshot .getId ()).setVer (1 ).build ();
157
+ when (range .toReseter (eq (newSnapshot ))).thenReturn (reseter );
155
158
CompletableFuture <Void > secondRestore = restorer .restoreFrom ("leader" , newSnapshot );
156
159
157
160
verify (reseter , times (1 )).abort ();
0 commit comments