From 50bd5842c3cbaa3a68f95cb84f6d893c7fabf50e Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <nhat.nguyen@elastic.co>
Date: Thu, 26 Dec 2019 08:56:00 -0500
Subject: [PATCH] Fix testCancelRecoveryDuringPhase1 (#50449)

testCancelRecoveryDuringPhase1 uses a mock of IndexShard, which can't
create retention leases. We need to stub method createRetentionLease.

Relates #50351
Closes #50424
---
 .../indices/recovery/RecoverySourceHandler.java      |  2 +-
 .../indices/recovery/RecoverySourceHandlerTests.java | 12 ++++++++++--
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/indices/recovery/RecoverySourceHandler.java b/server/src/main/java/org/elasticsearch/indices/recovery/RecoverySourceHandler.java
index 07db659299e7c..3e132c979d7b6 100644
--- a/server/src/main/java/org/elasticsearch/indices/recovery/RecoverySourceHandler.java
+++ b/server/src/main/java/org/elasticsearch/indices/recovery/RecoverySourceHandler.java
@@ -555,7 +555,7 @@ void phase1(IndexCommit snapshot, long startingSeqNo, IntSupplier translogOps, A
         }
     }
 
-    private void createRetentionLease(final long startingSeqNo, ActionListener<RetentionLease> listener) {
+    void createRetentionLease(final long startingSeqNo, ActionListener<RetentionLease> listener) {
         runUnderPrimaryPermit(() -> {
                 // Clone the peer recovery retention lease belonging to the source shard. We are retaining history between the the local
                 // checkpoint of the safe commit we're creating and this lease's retained seqno with the retention lock, and by cloning an
diff --git a/server/src/test/java/org/elasticsearch/indices/recovery/RecoverySourceHandlerTests.java b/server/src/test/java/org/elasticsearch/indices/recovery/RecoverySourceHandlerTests.java
index a10a8fc0410b6..999028cb083fd 100644
--- a/server/src/test/java/org/elasticsearch/indices/recovery/RecoverySourceHandlerTests.java
+++ b/server/src/test/java/org/elasticsearch/indices/recovery/RecoverySourceHandlerTests.java
@@ -64,6 +64,7 @@
 import org.elasticsearch.index.mapper.SeqNoFieldMapper;
 import org.elasticsearch.index.mapper.Uid;
 import org.elasticsearch.index.seqno.ReplicationTracker;
+import org.elasticsearch.index.seqno.RetentionLease;
 import org.elasticsearch.index.seqno.RetentionLeases;
 import org.elasticsearch.index.seqno.SeqNoStats;
 import org.elasticsearch.index.seqno.SequenceNumbers;
@@ -628,7 +629,6 @@ public void writeFileChunk(StoreFileMetaData md, long position, BytesReference c
         store.close();
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/50424")
     public void testCancelRecoveryDuringPhase1() throws Exception {
         Store store = newStore(createTempDir("source"), false);
         IndexShard shard = mock(IndexShard.class);
@@ -677,8 +677,16 @@ public void cleanFiles(int totalTranslogOps, long globalCheckpoint, Store.Metada
                 }
             }
         };
+        final StartRecoveryRequest startRecoveryRequest = getStartRecoveryRequest();
         final RecoverySourceHandler handler = new RecoverySourceHandler(
-            shard, recoveryTarget, threadPool, getStartRecoveryRequest(), between(1, 16), between(1, 4));
+            shard, recoveryTarget, threadPool, startRecoveryRequest, between(1, 16), between(1, 4)) {
+            @Override
+            void createRetentionLease(long startingSeqNo, ActionListener<RetentionLease> listener) {
+                final String leaseId = ReplicationTracker.getPeerRecoveryRetentionLeaseId(startRecoveryRequest.targetNode().getId());
+                listener.onResponse(new RetentionLease(leaseId, startingSeqNo, threadPool.absoluteTimeInMillis(),
+                    ReplicationTracker.PEER_RECOVERY_RETENTION_LEASE_SOURCE));
+            }
+        };
         cancelRecovery.set(() -> handler.cancel("test"));
         final StepListener<RecoverySourceHandler.SendFileResult> phase1Listener = new StepListener<>();
         try {