From e1fd53c718bb7b49076803c4e5059982c33b74af Mon Sep 17 00:00:00 2001
From: Aaron Train <aaron.train@gmail.com>
Date: Fri, 6 Dec 2019 14:53:18 -0500
Subject: [PATCH] Closes #6873: Support many more Content-Type in MockWebServer
 (UI tests) (#7001)

Ensure our MockResponse can handle different Content-Type e.g, images, to help support other tests
---
 .../mozilla/fenix/helpers/MockWebServer.kt    | 38 +++++++++++++++++--
 1 file changed, 34 insertions(+), 4 deletions(-)

diff --git a/app/src/androidTest/java/org/mozilla/fenix/helpers/MockWebServer.kt b/app/src/androidTest/java/org/mozilla/fenix/helpers/MockWebServer.kt
index c51a826b3..70a345b9f 100644
--- a/app/src/androidTest/java/org/mozilla/fenix/helpers/MockWebServer.kt
+++ b/app/src/androidTest/java/org/mozilla/fenix/helpers/MockWebServer.kt
@@ -12,8 +12,11 @@ import okhttp3.mockwebserver.Dispatcher
 import okhttp3.mockwebserver.MockResponse
 import okhttp3.mockwebserver.MockWebServer
 import okhttp3.mockwebserver.RecordedRequest
+import okio.Buffer
+import okio.source
 import org.mozilla.fenix.helpers.ext.toUri
 import java.io.IOException
+import java.io.InputStream
 
 object MockWebServerHelper {
 
@@ -37,7 +40,7 @@ object MockWebServerHelper {
  * If the dispatcher is unable to read a requested asset, it will fail the test by throwing an
  * Exception on the main thread.
  *
- * @sample [org.mozilla.tv.firefox.ui.BasicNavigationTest.basicNavigationTest]
+ * @sample [org.mozilla.fenix.ui.NavigationToolbarTest.visitURLTest]
  */
 const val HTTP_OK = 200
 const val HTTP_NOT_FOUND = 404
@@ -47,16 +50,43 @@ class AndroidAssetDispatcher : Dispatcher() {
 
     override fun dispatch(request: RecordedRequest): MockResponse {
         val assetManager = InstrumentationRegistry.getInstrumentation().context.assets
-        val assetContents = try {
+        try {
             val pathNoLeadingSlash = request.path.drop(1)
             assetManager.open(pathNoLeadingSlash).use { inputStream ->
-                inputStream.bufferedReader().use { it.readText() }
+                return fileToResponse(pathNoLeadingSlash, inputStream)
             }
         } catch (e: IOException) { // e.g. file not found.
             // We're on a background thread so we need to forward the exception to the main thread.
             mainThreadHandler.postAtFrontOfQueue { throw e }
             return MockResponse().setResponseCode(HTTP_NOT_FOUND)
         }
-        return MockResponse().setResponseCode(HTTP_OK).setBody(assetContents)
+    }
+}
+
+@Throws(IOException::class)
+private fun fileToResponse(path: String, file: InputStream): MockResponse {
+    return MockResponse()
+        .setResponseCode(HTTP_OK)
+        .setBody(fileToBytes(file))
+        .addHeader("content-type: " + contentType(path))
+}
+
+@Throws(IOException::class)
+private fun fileToBytes(file: InputStream): Buffer? {
+    val result = Buffer()
+    result.writeAll(file.source())
+    return result
+}
+
+private fun contentType(path: String): String? {
+    return when {
+        path.endsWith(".png") -> "image/png"
+        path.endsWith(".jpg") -> "image/jpeg"
+        path.endsWith(".jpeg") -> "image/jpeg"
+        path.endsWith(".gif") -> "image/gif"
+        path.endsWith(".svg") -> "image/svg+xml"
+        path.endsWith(".html") -> "text/html; charset=utf-8"
+        path.endsWith(".txt") -> "text/plain; charset=utf-8"
+        else -> "application/octet-stream"
     }
 }