Allow unstashing values into keys (#24685)

This is almost exclusively for docs test which frequently match the
entire response. This allow something like:
```
  - set: {nodes.$master.http.publish_address: host}
  - match:
      $body:
        {
          "nodes": {
            $host: {
              ... stuff in here ...
            }
          }
        }
```

This should make it possible for the docs tests to work with
unpredictable keys.

buildSrc/src/main/groovy/org/elasticsearch/gradle/doc/RestTestsFromSnippetsTask.groovy

@@ -167,6 +167,7 @@ public class RestTestsFromSnippetsTask extends SnippetsTask {
                  * warning every time. */
                 current.println("  - skip:")
                 current.println("      features: ")
+                current.println("        - stash_in_key")
                 current.println("        - warnings")
             }
             if (test.skipTest) {

test/framework/src/main/java/org/elasticsearch/test/rest/yaml/Features.java

@@ -39,6 +39,7 @@ public final class Features {
             "catch_unauthorized",
             "embedded_stash_key",
             "headers",
+            "stash_in_key",
             "stash_in_path",
             "warnings",
             "yaml"));

test/framework/src/main/java/org/elasticsearch/test/rest/yaml/Stash.java

@@ -26,9 +26,11 @@ import org.elasticsearch.common.xcontent.ToXContent;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.TreeMap;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -121,35 +123,46 @@ public class Stash implements ToXContent {
      * Goes recursively against each map entry and replaces any string value starting with "$" with its
      * corresponding value retrieved from the stash
      */
+    @SuppressWarnings("unchecked") // Safe because we check that all the map keys are string in unstashObject
     public Map<String, Object> replaceStashedValues(Map<String, Object> map) throws IOException {
-        Map<String, Object> copy = new HashMap<>(map);
+        return (Map<String, Object>) unstashObject(map);
-        unstashObject(copy);
-        return copy;
     }
 
-    @SuppressWarnings("unchecked")
+    private Object unstashObject(Object obj) throws IOException {
-    private void unstashObject(Object obj) throws IOException {
         if (obj instanceof List) {
-            List list = (List) obj;
+            List<?> list = (List<?>) obj;
-            for (int i = 0; i < list.size(); i++) {
+            List<Object> result = new ArrayList<>();
-                Object o = list.get(i);
+            for (Object o : list) {
                 if (containsStashedValue(o)) {
-                    list.set(i, getValue(o.toString()));
+                    result.add(getValue(o.toString()));
                 } else {
-                    unstashObject(o);
+                    result.add(unstashObject(o));
                 }
             }
+            return result;
         }
         if (obj instanceof Map) {
-            Map<String, Object> map = (Map) obj;
+            Map<?, ?> map = (Map<?, ?>) obj;
-            for (Map.Entry<String, Object> entry : map.entrySet()) {
+            Map<String, Object> result = new HashMap<>();
-                if (containsStashedValue(entry.getValue())) {
+            for (Map.Entry<?, ?> entry : map.entrySet()) {
-                    entry.setValue(getValue(entry.getValue().toString()));
+                String key = (String) entry.getKey();
+                Object value = entry.getValue();
+                if (containsStashedValue(key)) {
+                    key = getValue(key).toString();
+                }
+                if (containsStashedValue(value)) {
+                    value = getValue(value.toString());
                 } else {
-                    unstashObject(entry.getValue());
+                    value = unstashObject(value);
+                }
+                if (null != result.putIfAbsent(key, value)) {
+                    throw new IllegalArgumentException("Unstashing has caused a key conflict! The map is [" + result + "] and the key is ["
+                            + entry.getKey() + "] which unstashes to [" + key + "]");
                 }
             }
+            return result;
         }
+        return obj;
     }
 
     @Override

test/framework/src/test/java/org/elasticsearch/test/rest/yaml/StashTests.java

@@ -20,27 +20,101 @@
 package org.elasticsearch.test.rest.yaml;
 
 import org.elasticsearch.test.ESTestCase;
-import org.elasticsearch.test.rest.yaml.Stash;
 
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 
 import static java.util.Collections.singletonMap;
+import static org.hamcrest.Matchers.not;
+import static org.hamcrest.Matchers.sameInstance;
 
 public class StashTests extends ESTestCase {
-    public void testReplaceStashedValuesEmbeddedStashKey() throws IOException {
+    public void testReplaceStashedValuesStashKeyInMapValue() throws IOException {
         Stash stash = new Stash();
-        stash.stashValue("stashed", "bar");
 
         Map<String, Object> expected = new HashMap<>();
         expected.put("key", singletonMap("a", "foobar"));
         Map<String, Object> map = new HashMap<>();
         Map<String, Object> map2 = new HashMap<>();
+        if (randomBoolean()) {
+            stash.stashValue("stashed", "bar");
             map2.put("a", "foo${stashed}");
+        } else {
+            stash.stashValue("stashed", "foobar");
+            map2.put("a", "$stashed");
+        }
         map.put("key", map2);
 
         Map<String, Object> actual = stash.replaceStashedValues(map);
         assertEquals(expected, actual);
+        assertThat(actual, not(sameInstance(map)));
+    }
+
+    public void testReplaceStashedValuesStashKeyInMapKey() throws IOException {
+        Stash stash = new Stash();
+
+        Map<String, Object> expected = new HashMap<>();
+        expected.put("key", singletonMap("foobar", "a"));
+        Map<String, Object> map = new HashMap<>();
+        Map<String, Object> map2 = new HashMap<>();
+        if (randomBoolean()) {
+            stash.stashValue("stashed", "bar");
+            map2.put("foo${stashed}", "a");
+        } else {
+            stash.stashValue("stashed", "foobar");
+            map2.put("$stashed", "a");
+        }
+        map.put("key", map2);
+
+        Map<String, Object> actual = stash.replaceStashedValues(map);
+        assertEquals(expected, actual);
+        assertThat(actual, not(sameInstance(map)));
+    }
+
+    public void testReplaceStashedValuesStashKeyInMapKeyConflicts() throws IOException {
+        Stash stash = new Stash();
+
+        Map<String, Object> map = new HashMap<>();
+        Map<String, Object> map2 = new HashMap<>();
+        String key;
+        if (randomBoolean()) {
+            stash.stashValue("stashed", "bar");
+            key = "foo${stashed}";
+        } else {
+            stash.stashValue("stashed", "foobar");
+            key = "$stashed";
+        }
+        map2.put(key, "a");
+        map2.put("foobar", "whatever");
+        map.put("key", map2);
+
+        Exception e = expectThrows(IllegalArgumentException.class, () -> stash.replaceStashedValues(map));
+        assertEquals(e.getMessage(), "Unstashing has caused a key conflict! The map is [{foobar=whatever}] and the key is ["
+                            + key + "] which unstashes to [foobar]");
+    }
+
+
+    public void testReplaceStashedValuesStashKeyInList() throws IOException {
+        Stash stash = new Stash();
+        stash.stashValue("stashed", "bar");
+
+        Map<String, Object> expected = new HashMap<>();
+        expected.put("key", Arrays.asList("foot", "foobar", 1));
+        Map<String, Object> map = new HashMap<>();
+        Object value;
+        if (randomBoolean()) {
+            stash.stashValue("stashed", "bar");
+            value = "foo${stashed}";
+        } else {
+            stash.stashValue("stashed", "foobar");
+            value = "$stashed";
+        }
+        map.put("key", Arrays.asList("foot", value, 1));
+
+        Map<String, Object> actual = stash.replaceStashedValues(map);
+        assertEquals(expected, actual);
+        assertThat(actual, not(sameInstance(map)));
     }
 }