From 55e543a694441ca0999af5a4f8f0f95dfaff1dca Mon Sep 17 00:00:00 2001
From: Gordon Brown <gordon.brown@elastic.co>
Date: Fri, 3 Sep 2021 17:39:58 -0600
Subject: [PATCH] Fix randomization in MlNodeShutdownIT (#77281)

This commit fixes the randomization in MlNodeShutdownIT to ensure that
the required values for the selected shutdown type are all present.
---
 .../xpack/ml/integration/MlNodeShutdownIT.java        | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/x-pack/plugin/ml/src/internalClusterTest/java/org/elasticsearch/xpack/ml/integration/MlNodeShutdownIT.java b/x-pack/plugin/ml/src/internalClusterTest/java/org/elasticsearch/xpack/ml/integration/MlNodeShutdownIT.java
index 53af09dc7194f..d5ed5f2545b24 100644
--- a/x-pack/plugin/ml/src/internalClusterTest/java/org/elasticsearch/xpack/ml/integration/MlNodeShutdownIT.java
+++ b/x-pack/plugin/ml/src/internalClusterTest/java/org/elasticsearch/xpack/ml/integration/MlNodeShutdownIT.java
@@ -74,6 +74,8 @@ public void testJobsVacateShuttingDownNode() throws Exception {
         });
 
         // Call the shutdown API for the chosen node.
+        final SingleNodeShutdownMetadata.Type type = randomFrom(SingleNodeShutdownMetadata.Type.values());
+        final String targetNodeName = type == SingleNodeShutdownMetadata.Type.REPLACE ? randomAlphaOfLengthBetween(10, 20) : null;
         client().execute(
             PutShutdownNodeAction.INSTANCE,
             new PutShutdownNodeAction.Request(
@@ -81,7 +83,7 @@ public void testJobsVacateShuttingDownNode() throws Exception {
                 randomFrom(SingleNodeShutdownMetadata.Type.values()),
                 "just testing",
                 null,
-                null)
+                targetNodeName)
         ).actionGet();
 
         // Wait for the desired end state of all 6 jobs running on nodes that are not shutting down.
@@ -144,14 +146,15 @@ public void testCloseJobVacatingShuttingDownNode() throws Exception {
         });
 
         // Call the shutdown API for the chosen node.
+        final SingleNodeShutdownMetadata.Type type = randomFrom(SingleNodeShutdownMetadata.Type.values());
+        final String targetNodeName = type == SingleNodeShutdownMetadata.Type.REPLACE ? randomAlphaOfLengthBetween(10, 20) : null;
         client().execute(
             PutShutdownNodeAction.INSTANCE,
             new PutShutdownNodeAction.Request(
-                nodeIdToShutdown.get(),
-                randomFrom(SingleNodeShutdownMetadata.Type.values()),
+                nodeIdToShutdown.get(), type,
                 "just testing",
                 null,
-                null)
+                targetNodeName)
         )
             .actionGet();