Supports to add a response for the given headers

Currently TestHttpServer only supports to add a response
for Location header. Make it to add a response for the
given headers so that any additional headers that should
be returned along with the response.

Bug: 181745786
Test: build
Original-Change: https://android-review.googlesource.com/1724310
Merged-In: I80b2fe43016b9cd12fc3cb95ee76f1ec9568ddfc
Change-Id: I80b2fe43016b9cd12fc3cb95ee76f1ec9568ddfc

1 files changed, 6 insertions, 4 deletions

diff --git a/common/testutils/devicetests/com/android/testutils/TestHttpServer.kt b/common/testutils/devicetests/com/android/testutils/TestHttpServer.kt
index 7aae8e3b..39ce4872 100644
--- a/common/testutils/devicetests/com/android/testutils/TestHttpServer.kt
+++ b/common/testutils/devicetests/com/android/testutils/TestHttpServer.kt
@@ -54,12 +54,12 @@ class TestHttpServer(host: String? = null) : NanoHTTPD(host, 0 /* auto-select th
     fun addResponse(
         uri: Uri,
         statusCode: Response.IStatus,
-        locationHeader: String? = null,
+        headers: Map<String, String>? = null,
         content: String = ""
     ) {
         addResponse(Request(uri.path
                 ?: "", Method.GET, uri.query ?: ""),
-                statusCode, locationHeader, content)
+                statusCode, headers, content)
     }
 
     /**
@@ -68,11 +68,13 @@ class TestHttpServer(host: String? = null) : NanoHTTPD(host, 0 /* auto-select th
     fun addResponse(
         request: Request,
         statusCode: Response.IStatus,
-        locationHeader: String? = null,
+        headers: Map<String, String>? = null,
         content: String = ""
     ) {
         val response = newFixedLengthResponse(statusCode, "text/plain", content)
-        locationHeader?.let { response.addHeader("Location", it) }
+        headers?.forEach {
+            (key, value) -> response.addHeader(key, value)
+        }
         responses[request] = response
     }