Add tests for setting custom headers.

Test: atest CtsNetHttpTestCases

Change-Id: I7085c2e29a3bf6a467c2db1112e2df1fad8efb12
diff --git a/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java b/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
index 2ec035b..a4dca19 100644
--- a/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
+++ b/Cronet/tests/cts/src/android/net/http/cts/UrlRequestTest.java
@@ -29,6 +29,7 @@
 import static org.junit.Assert.assertTrue;
 
 import android.content.Context;
+import android.net.http.HeaderBlock;
 import android.net.http.HttpEngine;
 import android.net.http.HttpException;
 import android.net.http.InlineExecutionProhibitedException;
@@ -43,6 +44,7 @@
 import android.net.http.cts.util.TestUrlRequestCallback;
 import android.net.http.cts.util.TestUrlRequestCallback.ResponseStep;
 import android.net.http.cts.util.UploadDataProviders;
+import android.webkit.cts.CtsTestServer;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -59,11 +61,16 @@
 import java.net.URLEncoder;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 @RunWith(AndroidJUnit4.class)
 public class UrlRequestTest {
@@ -312,6 +319,50 @@
         assertThat(mCallback.mResponseAsString).isEqualTo(body);
     }
 
+    @Test
+    public void testUrlRequest_customHeaders() throws Exception {
+        UrlRequest.Builder builder = createUrlRequestBuilder(mTestServer.getEchoHeadersUrl());
+
+        List<Map.Entry<String, String>> expectedHeaders = Arrays.asList(
+                Map.entry("Authorization", "Basic QWxhZGRpbjpvcGVuIHNlc2FtZQ=="),
+                Map.entry("Max-Forwards", "10"),
+                Map.entry("X-Client-Data", "random custom header content"));
+
+        for (Map.Entry<String, String> header : expectedHeaders) {
+            builder.addHeader(header.getKey(), header.getValue());
+        }
+
+        builder.build().start();
+        mCallback.expectCallback(ResponseStep.ON_SUCCEEDED);
+
+        assertOKStatusCode(mCallback.mResponseInfo);
+
+        List<Map.Entry<String, String>> echoedHeaders =
+                extractEchoedHeaders(mCallback.mResponseInfo.getHeaders());
+
+        // The implementation might decide to add more headers like accepted encodings it handles
+        // internally so the server is likely to see more headers than explicitly set
+        // by the developer.
+        assertThat(echoedHeaders)
+                .containsAtLeastElementsIn(expectedHeaders);
+    }
+
+    private static List<Map.Entry<String, String>> extractEchoedHeaders(HeaderBlock headers) {
+        return headers.getAsList()
+                .stream()
+                .flatMap(input -> {
+                    if (input.getKey().startsWith(CtsTestServer.ECHOED_RESPONSE_HEADER_PREFIX)) {
+                        String strippedKey =
+                                input.getKey().substring(
+                                        CtsTestServer.ECHOED_RESPONSE_HEADER_PREFIX.length());
+                        return Stream.of(Map.entry(strippedKey, input.getValue()));
+                    } else {
+                        return Stream.empty();
+                    }
+                })
+                .collect(Collectors.toList());
+    }
+
     private static class StubUrlRequestCallback implements UrlRequest.Callback {
 
         @Override