Implement the ServiceEntitlementException

Bug: 177544547
Change-Id: I284e9be6f040c29974abd3e8d35df08a6147399c
Merged-In: I284e9be6f040c29974abd3e8d35df08a6147399c
Test: EapAkaApiTest, EapAkaResponseTest, HttpClientTest
(cherry picked from commit 13a767f5dfd038f9b099698da8aaedb1707eb45a)

11 files changed, 348 insertions, 198 deletions

diff --git a/java/com/android/libraries/entitlement/ServiceEntitlement.java b/java/com/android/libraries/entitlement/ServiceEntitlement.java
index be50605..f788603 100644
--- a/java/com/android/libraries/entitlement/ServiceEntitlement.java
+++ b/java/com/android/libraries/entitlement/ServiceEntitlement.java
@@ -19,11 +19,10 @@ package com.android.libraries.entitlement;
 import android.content.Context;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.libraries.entitlement.eapaka.EapAkaApi;
 
-import com.google.common.annotations.VisibleForTesting;
-
 import java.util.List;
 
 /**
diff --git a/java/com/android/libraries/entitlement/ServiceEntitlementException.java b/java/com/android/libraries/entitlement/ServiceEntitlementException.java
index fa476ce..d650726 100644
--- a/java/com/android/libraries/entitlement/ServiceEntitlementException.java
+++ b/java/com/android/libraries/entitlement/ServiceEntitlementException.java
@@ -46,29 +46,71 @@ public class ServiceEntitlementException extends Exception {
      */
     public static final int ERROR_HTTP_STATUS_NOT_SUCCESS = 4;
 
-    public ServiceEntitlementException(String message) {
-        // TODO(b/177544547): add implementation
+    /**
+     * HTTP response received with a malformed format. e.g. the response with content-type JSON but
+     * failing JSON parser.
+     */
+    public static final int MALFORMED_HTTP_RESPONSE = 5;
+
+    /**
+     * Default HTTP status if not been specified.
+     */
+    private static final int HTTP_STATUS_UNSPECIFIED = 0;
+
+    /**
+     * An empty string if Retry-After header in HTTP response not been specified.
+     */
+    private static final String RETRY_AFTER_UNSPECIFIED = "";
+
+    private int mError;
+    private int mHttpStatus;
+    private String mRetryAfter;
+
+    public ServiceEntitlementException(int error, String message) {
+        this(error, HTTP_STATUS_UNSPECIFIED, RETRY_AFTER_UNSPECIFIED, message);
+    }
+
+    public ServiceEntitlementException(int error, int httpStatus, String message) {
+        this(error, httpStatus, RETRY_AFTER_UNSPECIFIED, message);
+    }
+
+    public ServiceEntitlementException(
+            int error, int httpStatus, String retryAfter, String message) {
+        super(message);
+        this.mError = error;
+        this.mHttpStatus = httpStatus;
+        this.mRetryAfter = retryAfter;
+    }
+
+    public ServiceEntitlementException(int error, String message, Throwable cause) {
+        this(error, HTTP_STATUS_UNSPECIFIED, RETRY_AFTER_UNSPECIFIED, message, cause);
+    }
+
+    public ServiceEntitlementException(int error, int httpStatus, String message, Throwable cause) {
+        this(error, httpStatus, RETRY_AFTER_UNSPECIFIED, message, cause);
     }
 
     public ServiceEntitlementException(
             int error, int httpStatus, String retryAfter, String message, Throwable cause) {
-        // TODO(b/177544547): add implementation
+        super(message, cause);
+        this.mError = error;
+        this.mHttpStatus = httpStatus;
+        this.mRetryAfter = retryAfter;
     }
 
     /**
-     * Returns the error code, see {@link #ERROR_*}.
+     * Returns the error code, see {@link #ERROR_*}. {@link #ERROR_UNKNOWN} if not been specified.
      */
     public int getErrorCode() {
-        // TODO(b/177544547): add implementation
-        return ERROR_UNKNOWN;
+        return mError;
     }
 
     /**
-     * Returns the HTTP status code returned by entitlement server; 0 if unavailable.
+     * Returns the HTTP status code returned by entitlement server; {@link #HTTP_STATUS_UNSPECIFIED}
+     * if not been specified.
      */
     public int getHttpStatus() {
-        // TODO(b/177544547): add implementation
-        return ERROR_SEVER_NOT_CONNECTABLE;
+        return mHttpStatus;
     }
 
     /**
@@ -79,7 +121,6 @@ public class ServiceEntitlementException extends Exception {
      * https://tools.ietf.org/html/rfc7231#section-7.1.3
      */
     public String getRetryAfter() {
-        // TODO(b/177544547): add implementation
-        return null;
+        return mRetryAfter;
     }
 }
diff --git a/java/com/android/libraries/entitlement/eapaka/EapAkaApi.java b/java/com/android/libraries/entitlement/eapaka/EapAkaApi.java
index a145ebc..9b450b6 100644
--- a/java/com/android/libraries/entitlement/eapaka/EapAkaApi.java
+++ b/java/com/android/libraries/entitlement/eapaka/EapAkaApi.java
@@ -16,6 +16,8 @@
 
 package com.android.libraries.entitlement.eapaka;
 
+import static com.android.libraries.entitlement.ServiceEntitlementException.MALFORMED_HTTP_RESPONSE;
+
 import android.content.Context;
 import android.net.Uri;
 import android.telephony.TelephonyManager;
@@ -23,6 +25,7 @@ import android.text.TextUtils;
 import android.util.Log;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.libraries.entitlement.ServiceEntitlementException;
 import com.android.libraries.entitlement.ServiceEntitlementRequest;
@@ -32,7 +35,6 @@ import com.android.libraries.entitlement.http.HttpConstants.RequestMethod;
 import com.android.libraries.entitlement.http.HttpRequest;
 import com.android.libraries.entitlement.http.HttpResponse;
 
-import com.google.common.annotations.VisibleForTesting;
 import com.google.common.net.HttpHeaders;
 
 import org.json.JSONException;
@@ -98,6 +100,16 @@ public class EapAkaApi {
      * Attribute name of the app name.
      */
     private static final String APP_NAME = "app_name";
+    /**
+     * Expected content type of the response body.
+     */
+    private static final String ACCEPT_CONTENT_TYPE_JSON_AND_XML =
+            "application/vnd.gsma.eap-relay.v1.0+json, text/vnd.wap.connectivity-xml";
+    /**
+     * Required content type to the response body.
+     */
+    private static final String REQUEST_CONTENT_TYPE_JSON =
+            "application/vnd.gsma.eap-relay.v1.0+json";
 
     private final Context mContext;
     private final int mSimSubscriptionId;
@@ -132,64 +144,57 @@ public class EapAkaApi {
                 HttpRequest.builder()
                         .setUrl(entitlementStatusUrl(appId, serverUrl, request))
                         .setRequestMethod(RequestMethod.GET)
-                        .addRequestProperty(
-                                HttpHeaders.ACCEPT,
-                                "application/vnd.gsma.eap-relay.v1.0+json, text/vnd.wap"
-                                        + ".connectivity-xml")
+                        .addRequestProperty(HttpHeaders.ACCEPT, ACCEPT_CONTENT_TYPE_JSON_AND_XML)
                         .build();
         HttpResponse response = mHttpClient.request(httpRequest);
-        if (response == null) {
-            throw new ServiceEntitlementException("Null http response");
-        }
-        if (response.contentType() == ContentType.JSON) {
+        if (request.authenticationToken().isEmpty()) {
+            // EapAka token challenge for initial AuthN
+            if (response.contentType() != ContentType.JSON) {
+                throw new ServiceEntitlementException(
+                        MALFORMED_HTTP_RESPONSE, "Unexpected http ContentType");
+            }
+
+            Log.d(TAG, "initial AuthN");
+            String responseData = "";
             try {
-                // EapAka token challenge for initial AuthN
-                Log.d(TAG, "initial AuthN");
-                String akaChallengeResponse =
-                        new EapAkaResponse(
-                                new JSONObject(response.body()).getString(EAP_CHALLENGE_RESPONSE))
-                                .getEapAkaChallengeResponse(mContext, mSimSubscriptionId);
-                JSONObject postData = new JSONObject();
-                postData.put(EAP_CHALLENGE_RESPONSE, akaChallengeResponse);
-                return challengeResponse(postData, serverUrl);
+                responseData = new JSONObject(response.body()).getString(
+                        EAP_CHALLENGE_RESPONSE);
             } catch (JSONException jsonException) {
-                Log.e(TAG, "queryEntitlementStatus failed. jsonException: " + jsonException);
-                return null;
+                throw new ServiceEntitlementException(
+                        MALFORMED_HTTP_RESPONSE, "Failed to parse json object", jsonException);
             }
-        } else if (response.contentType() == ContentType.XML) {
+            return challengeResponse(
+                    new EapAkaResponse(responseData).getEapAkaChallengeResponse(mContext,
+                            mSimSubscriptionId), serverUrl);
+        } else {
             // Result of fast AuthN
             Log.d(TAG, "fast AuthN");
             return response.body();
         }
-        throw new ServiceEntitlementException("Unexpected http ContentType");
     }
 
-    private String challengeResponse(JSONObject postData, String serverUrl)
+    private String challengeResponse(String akaChallengeResponse, String serverUrl)
             throws ServiceEntitlementException {
         Log.d(TAG, "challengeResponse");
+        JSONObject postData = new JSONObject();
+        try {
+            postData.put(EAP_CHALLENGE_RESPONSE, akaChallengeResponse);
+        } catch (JSONException jsonException) {
+            throw new ServiceEntitlementException(
+                    MALFORMED_HTTP_RESPONSE, "Failed to put post data", jsonException);
+        }
         HttpRequest request =
                 HttpRequest.builder()
                         .setUrl(serverUrl)
                         .setRequestMethod(RequestMethod.POST)
                         .setPostData(postData)
-                        .addRequestProperty(
-                                HttpHeaders.ACCEPT,
-                                "application/vnd.gsma.eap-relay.v1.0+json, text/vnd.wap"
-                                        + ".connectivity-xml")
-                        .addRequestProperty(HttpHeaders.CONTENT_TYPE,
-                                "application/vnd.gsma.eap-relay.v1.0+json")
+                        .addRequestProperty(HttpHeaders.ACCEPT, ACCEPT_CONTENT_TYPE_JSON_AND_XML)
+                        .addRequestProperty(HttpHeaders.CONTENT_TYPE, REQUEST_CONTENT_TYPE_JSON)
                         .build();
-
-        HttpResponse response = mHttpClient.request(request);
-        if (response == null || response.contentType() != ContentType.XML) {
-            throw new ServiceEntitlementException("Unexpected http response.");
-        }
-
-        return response.body();
+        return mHttpClient.request(request).body();
     }
 
-    @VisibleForTesting
-    String entitlementStatusUrl(
+    private String entitlementStatusUrl(
             String appId, String serverUrl, ServiceEntitlementRequest request) {
         TelephonyManager telephonyManager = mContext.getSystemService(
                 TelephonyManager.class).createForSubscriptionId(mSimSubscriptionId);
diff --git a/java/com/android/libraries/entitlement/eapaka/EapAkaResponse.java b/java/com/android/libraries/entitlement/eapaka/EapAkaResponse.java
index 1d9be2b..62667b0 100644
--- a/java/com/android/libraries/entitlement/eapaka/EapAkaResponse.java
+++ b/java/com/android/libraries/entitlement/eapaka/EapAkaResponse.java
@@ -16,6 +16,8 @@
 
 package com.android.libraries.entitlement.eapaka;
 
+import static com.android.libraries.entitlement.ServiceEntitlementException.ERROR_ICC_AUTHENTICATION_NOT_AVAILABLE;
+
 import android.content.Context;
 import android.telephony.TelephonyManager;
 import android.text.TextUtils;
@@ -23,12 +25,11 @@ import android.util.Base64;
 import android.util.Log;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.libraries.entitlement.ServiceEntitlementException;
 import com.android.libraries.entitlement.eapaka.utils.BytesConverter;
 
-import com.google.common.annotations.VisibleForTesting;
-
 import java.security.InvalidKeyException;
 import java.security.NoSuchAlgorithmException;
 
@@ -216,7 +217,8 @@ class EapAkaResponse {
     public String getEapAkaChallengeResponse(Context context, int simSubscriptionId)
             throws ServiceEntitlementException {
         if (!mValid) {
-            throw new ServiceEntitlementException("EAP-AKA Challenge message not valid!");
+            throw new ServiceEntitlementException(
+                    ERROR_ICC_AUTHENTICATION_NOT_AVAILABLE, "EAP-AKA Challenge message not valid!");
         }
 
         TelephonyManager telephonyManager =
@@ -241,7 +243,8 @@ class EapAkaResponse {
                         securityContext.getCk());
         // K_aut is the key used to calculate MAC
         if (mk.getAut() == null) {
-            throw new ServiceEntitlementException("Can't generate K_Aut!");
+            throw new ServiceEntitlementException(
+                    ERROR_ICC_AUTHENTICATION_NOT_AVAILABLE, "Can't generate K_Aut!");
         }
 
         // generate EAP-AKA Challenge Response message
@@ -249,6 +252,7 @@ class EapAkaResponse {
                 generateEapAkaChallengeResponse(securityContext.getRes(), mk.getAut());
         if (challengeResponse == null) {
             throw new ServiceEntitlementException(
+                    ERROR_ICC_AUTHENTICATION_NOT_AVAILABLE,
                     "Failed to generate EAP-AKA Challenge Response data!");
         }
 
diff --git a/java/com/android/libraries/entitlement/eapaka/EapAkaSecurityContext.java b/java/com/android/libraries/entitlement/eapaka/EapAkaSecurityContext.java
index 925e258..cdb9427 100644
--- a/java/com/android/libraries/entitlement/eapaka/EapAkaSecurityContext.java
+++ b/java/com/android/libraries/entitlement/eapaka/EapAkaSecurityContext.java
@@ -16,6 +16,8 @@
 
 package com.android.libraries.entitlement.eapaka;
 
+import static com.android.libraries.entitlement.ServiceEntitlementException.ERROR_ICC_AUTHENTICATION_NOT_AVAILABLE;
+
 import android.text.TextUtils;
 import android.util.Base64;
 import android.util.Log;
@@ -52,7 +54,9 @@ class EapAkaSecurityContext {
         EapAkaSecurityContext securityContext = new EapAkaSecurityContext();
         securityContext.parseResponseData(response);
         if (!securityContext.isValid()) {
-            throw new ServiceEntitlementException("Invalid SIM EAP-AKA authentication response!");
+            throw new ServiceEntitlementException(
+                ERROR_ICC_AUTHENTICATION_NOT_AVAILABLE,
+                "Invalid SIM EAP-AKA authentication response!");
         }
         return securityContext;
     }
diff --git a/java/com/android/libraries/entitlement/http/HttpClient.java b/java/com/android/libraries/entitlement/http/HttpClient.java
index 9d102ce..0facae8 100644
--- a/java/com/android/libraries/entitlement/http/HttpClient.java
+++ b/java/com/android/libraries/entitlement/http/HttpClient.java
@@ -16,6 +16,9 @@
 
 package com.android.libraries.entitlement.http;
 
+import static com.android.libraries.entitlement.ServiceEntitlementException.ERROR_HTTP_STATUS_NOT_SUCCESS;
+import static com.android.libraries.entitlement.ServiceEntitlementException.ERROR_SEVER_NOT_CONNECTABLE;
+import static com.android.libraries.entitlement.ServiceEntitlementException.MALFORMED_HTTP_RESPONSE;
 import static com.android.libraries.entitlement.http.HttpConstants.RequestMethod.POST;
 
 import static com.google.common.base.Strings.nullToEmpty;
@@ -32,6 +35,8 @@ import com.android.libraries.entitlement.ServiceEntitlementException;
 import com.android.libraries.entitlement.http.HttpConstants.ContentType;
 import com.android.libraries.entitlement.utils.StreamUtils;
 
+import com.google.common.net.HttpHeaders;
+
 import java.io.DataOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
@@ -56,14 +61,10 @@ public class HttpClient {
     @WorkerThread
     // TODO(b/177544547): Add debug messages
     public HttpResponse request(HttpRequest request) throws ServiceEntitlementException {
+        logd("HttpClient.request url: " + request.url());
+        createConnection(request);
+        logd("HttpClient.request headers (partial): " + mConnection.getRequestProperties());
         try {
-            logd("HttpClient.request url: " + request.url());
-            createConnection(request);
-            if (mConnection == null) {
-                logd("HttpClient.request connection is null");
-                throw new ServiceEntitlementException("No connection");
-            }
-            logd("HttpClient.request headers (partial): " + mConnection.getRequestProperties());
             if (POST.equals(request.requestMethod())) {
                 try (OutputStream out = new DataOutputStream(mConnection.getOutputStream())) {
                     out.write(request.postData().toString().getBytes(UTF_8));
@@ -74,13 +75,11 @@ public class HttpClient {
             HttpResponse response = getHttpResponse(mConnection);
             Log.d(TAG, "HttpClient.response : " + response);
             return response;
-        } catch (IOException e) {
-            InputStream errorStream = mConnection.getErrorStream();
-            Log.e(
-                    TAG,
-                    "HttpClient.request() error: " + StreamUtils.inputStreamToStringSafe(
-                            errorStream));
-            throw new ServiceEntitlementException("request failed! exception: " + e.getMessage());
+        } catch (IOException ioe) {
+            throw new ServiceEntitlementException(
+                    ERROR_HTTP_STATUS_NOT_SUCCESS,
+                    StreamUtils.inputStreamToStringSafe(mConnection.getErrorStream()),
+                    ioe);
         } finally {
             closeConnection();
         }
@@ -103,9 +102,9 @@ public class HttpClient {
             if (POST.equals(request.requestMethod())) {
                 mConnection.setDoOutput(true);
             }
-        } catch (IOException e) {
-            Log.e(TAG, "IOException: " + e.getMessage());
-            throw new ServiceEntitlementException("Configure connection failed!" + e.getMessage());
+        } catch (IOException ioe) {
+            throw new ServiceEntitlementException(
+                    ERROR_SEVER_NOT_CONNECTABLE, "Configure connection failed!", ioe);
         }
     }
 
@@ -118,28 +117,31 @@ public class HttpClient {
 
     private static HttpResponse getHttpResponse(HttpURLConnection connection)
             throws ServiceEntitlementException {
+        HttpResponse.Builder responseBuilder = HttpResponse.builder();
+        responseBuilder.setContentType(getContentType(connection));
         try {
             int responseCode = connection.getResponseCode();
             logd("HttpClient.response headers: " + connection.getHeaderFields());
             if (responseCode != HttpURLConnection.HTTP_OK) {
-                throw new ServiceEntitlementException(
-                        ServiceEntitlementException.ERROR_HTTP_STATUS_NOT_SUCCESS, responseCode,
-                        null,
-                        "Invalid connection response", null);
+                throw new ServiceEntitlementException(ERROR_HTTP_STATUS_NOT_SUCCESS, responseCode,
+                        connection.getHeaderField(HttpHeaders.RETRY_AFTER),
+                        "Invalid connection response");
             }
+            responseBuilder.setResponseCode(responseCode);
+            responseBuilder.setResponseMessage(nullToEmpty(connection.getResponseMessage()));
+        } catch (IOException e) {
+            throw new ServiceEntitlementException(
+                    ERROR_HTTP_STATUS_NOT_SUCCESS, "Read response code failed!", e);
+        }
+        try {
             String responseBody = readResponse(connection);
             logd("HttpClient.response body: " + responseBody);
-            return HttpResponse.builder()
-                    .setContentType(getContentType(connection))
-                    .setBody(responseBody)
-                    .setResponseCode(responseCode)
-                    .setResponseMessage(nullToEmpty(connection.getResponseMessage()))
-                    .build();
+            responseBuilder.setBody(responseBody);
         } catch (IOException e) {
             throw new ServiceEntitlementException(
-                    ServiceEntitlementException.ERROR_HTTP_STATUS_NOT_SUCCESS, 0, null,
-                    "Read response failed!", e);
+            MALFORMED_HTTP_RESPONSE, "Read response body/message failed!", e);
         }
+        return responseBuilder.build();
     }
 
     private static String readResponse(URLConnection connection) throws IOException {
diff --git a/tests/src/com/android/libraries/entitlement/eapaka/EapAkaApiTest.java b/tests/src/com/android/libraries/entitlement/eapaka/EapAkaApiTest.java
index dde3d06..7a79526 100644
--- a/tests/src/com/android/libraries/entitlement/eapaka/EapAkaApiTest.java
+++ b/tests/src/com/android/libraries/entitlement/eapaka/EapAkaApiTest.java
@@ -21,6 +21,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
+import static org.testng.Assert.fail;
 
 import android.content.Context;
 import android.telephony.TelephonyManager;
@@ -29,11 +30,13 @@ import androidx.test.core.app.ApplicationProvider;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.libraries.entitlement.ServiceEntitlement;
+import com.android.libraries.entitlement.ServiceEntitlementException;
 import com.android.libraries.entitlement.ServiceEntitlementRequest;
 import com.android.libraries.entitlement.http.HttpClient;
 import com.android.libraries.entitlement.http.HttpConstants.ContentType;
 import com.android.libraries.entitlement.http.HttpResponse;
 
+import org.json.JSONException;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -140,4 +143,44 @@ public class EapAkaApiTest {
                 mEapAkaApi.queryEntitlementStatus(ServiceEntitlement.APP_VOWIFI, TEST_URL, request))
                 .isEqualTo(RESPONSE_XML);
     }
+
+    @Test
+    public void queryEntitlementStatus_noAuthenticationTokenContentTypeNotJson_throwException()
+            throws Exception {
+        HttpResponse xmlResponse =
+                HttpResponse.builder().setContentType(ContentType.XML).setBody(RESPONSE_XML)
+                        .build();
+        when(mMockHttpClient.request(any())).thenReturn(xmlResponse);
+
+        ServiceEntitlementRequest request = ServiceEntitlementRequest.builder().build();
+
+        try {
+            mEapAkaApi.queryEntitlementStatus(ServiceEntitlement.APP_VOWIFI, TEST_URL, request);
+            fail();
+        } catch (ServiceEntitlementException exception) {
+            assertThat(exception.getErrorCode()).isEqualTo(
+                    ServiceEntitlementException.MALFORMED_HTTP_RESPONSE);
+            assertThat(exception.getMessage()).isEqualTo("Unexpected http ContentType");
+        }
+    }
+
+    @Test
+    public void queryEntitlementStatus_noAuthenticationTokenEmptyResponseBody_throwException()
+            throws Exception {
+        HttpResponse eapChallengeResponse =
+                HttpResponse.builder().setContentType(ContentType.JSON).build();
+        when(mMockHttpClient.request(any())).thenReturn(eapChallengeResponse);
+
+        ServiceEntitlementRequest request = ServiceEntitlementRequest.builder().build();
+
+        try {
+            mEapAkaApi.queryEntitlementStatus(ServiceEntitlement.APP_VOWIFI, TEST_URL, request);
+            fail();
+        } catch (ServiceEntitlementException exception) {
+            assertThat(exception.getErrorCode()).isEqualTo(
+                    ServiceEntitlementException.MALFORMED_HTTP_RESPONSE);
+            assertThat(exception.getMessage()).isEqualTo("Failed to parse json object");
+            assertThat(exception.getCause()).isInstanceOf(JSONException.class);
+        }
+    }
 }
\ No newline at end of file
diff --git a/tests/src/com/android/libraries/entitlement/eapaka/EapAkaResponseTest.java b/tests/src/com/android/libraries/entitlement/eapaka/EapAkaResponseTest.java
index 652cba6..068346a 100644
--- a/tests/src/com/android/libraries/entitlement/eapaka/EapAkaResponseTest.java
+++ b/tests/src/com/android/libraries/entitlement/eapaka/EapAkaResponseTest.java
@@ -19,6 +19,7 @@ package com.android.libraries.entitlement.eapaka;
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.fail;
 
 import android.content.Context;
 import android.util.Base64;
@@ -253,6 +254,23 @@ public class EapAkaResponseTest {
     }
 
     @Test
+    public void parseEapAkaChallengeRequest_notValid_throwException() {
+        byte[] data = convertHexStringToBytes(EAP_AKA_CHALLENGE_REQUEST_WITH_WRONG_LENGTH);
+        String encodedData = Base64.encodeToString(data, Base64.NO_WRAP).trim();
+
+        EapAkaResponse message = new EapAkaResponse(encodedData);
+
+        try {
+            message.getEapAkaChallengeResponse(mMockContext, SUB_ID);
+            fail();
+        } catch (ServiceEntitlementException exception) {
+            assertThat(exception.getErrorCode()).isEqualTo(
+                    ServiceEntitlementException.ERROR_ICC_AUTHENTICATION_NOT_AVAILABLE);
+            assertThat(exception.getMessage()).isEqualTo("EAP-AKA Challenge message not valid!");
+        }
+    }
+
+    @Test
     public void generateEapAkaChallengeResponse_pass() {
         EapAkaResponse message = new EapAkaResponse(EAP_AKA_CHALLENGE_REQUEST_EXPECTED);
 
diff --git a/tests/src/com/android/libraries/entitlement/eapaka/EapAkaSecurityContextTest.java b/tests/src/com/android/libraries/entitlement/eapaka/EapAkaSecurityContextTest.java
index 6d96abf..15c4c77 100644
--- a/tests/src/com/android/libraries/entitlement/eapaka/EapAkaSecurityContextTest.java
+++ b/tests/src/com/android/libraries/entitlement/eapaka/EapAkaSecurityContextTest.java
@@ -19,6 +19,7 @@ package com.android.libraries.entitlement.eapaka;
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.fail;
 
 import android.util.Base64;
 
@@ -90,8 +91,15 @@ public class EapAkaSecurityContextTest {
         byte[] data = convertHexStringToBytes(GSM_SECURITY_CONTEXT_RESPONSE_TAG_DC);
         String encodedData = Base64.encodeToString(data, Base64.NO_WRAP).trim();
 
-        assertThrows(
-                ServiceEntitlementException.class, () -> EapAkaSecurityContext.from(encodedData));
+        try {
+            EapAkaSecurityContext.from(encodedData);
+            fail();
+        } catch (ServiceEntitlementException exception) {
+            assertThat(exception.getErrorCode()).isEqualTo(
+                    ServiceEntitlementException.ERROR_ICC_AUTHENTICATION_NOT_AVAILABLE);
+            assertThat(exception.getMessage()).isEqualTo(
+                    "Invalid SIM EAP-AKA authentication response!");
+        }
     }
 
     @Test
diff --git a/tests/src/com/android/libraries/entitlement/http/HttpClientTest.java b/tests/src/com/android/libraries/entitlement/http/HttpClientTest.java
index 180ce09..3be93ca 100644
--- a/tests/src/com/android/libraries/entitlement/http/HttpClientTest.java
+++ b/tests/src/com/android/libraries/entitlement/http/HttpClientTest.java
@@ -18,7 +18,7 @@ package com.android.libraries.entitlement.http;
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.fail;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 
@@ -45,6 +45,7 @@ public class HttpClientTest {
     private static final String TEST_URL = "https://test.url";
     private static final String TEST_RESPONSE_BODY = "TEST_RESPONSE_BODY";
     private static final String CONTENT_TYPE_STRING_JSON = "application/json";
+    private static final String RETRY_AFTER = "RETRY_AFTER";
 
     private static FakeURLStreamHandler sFakeURLStreamHandler;
 
@@ -85,12 +86,22 @@ public class HttpClientTest {
                         .setResponseLocation(null)
                         .setResponseBody(TEST_RESPONSE_BODY.getBytes(UTF_8))
                         .setContentType(CONTENT_TYPE_STRING_JSON)
+                        .setRetryAfter(RETRY_AFTER)
                         .build();
         HttpRequest request =
                 HttpRequest.builder().setUrl(TEST_URL).setRequestMethod(RequestMethod.GET).build();
         Map<String, FakeResponse> response = ImmutableMap.of(TEST_URL, responseContent);
         sFakeURLStreamHandler.stubResponse(response);
 
-        assertThrows(ServiceEntitlementException.class, () -> mHttpClient.request(request));
+        try {
+            mHttpClient.request(request);
+            fail();
+        } catch (ServiceEntitlementException exception) {
+            assertThat(exception.getErrorCode()).isEqualTo(
+                    ServiceEntitlementException.ERROR_HTTP_STATUS_NOT_SUCCESS);
+            assertThat(exception.getHttpStatus()).isEqualTo(HttpURLConnection.HTTP_BAD_REQUEST);
+            assertThat(exception.getMessage()).isEqualTo("Invalid connection response");
+            assertThat(exception.getRetryAfter()).isEqualTo(RETRY_AFTER);
+        }
     }
 }
\ No newline at end of file
diff --git a/tests/utils/com/android/libraries/entitlement/testing/FakeURLStreamHandler.java b/tests/utils/com/android/libraries/entitlement/testing/FakeURLStreamHandler.java
index 0763747..d323412 100644
--- a/tests/utils/com/android/libraries/entitlement/testing/FakeURLStreamHandler.java
+++ b/tests/utils/com/android/libraries/entitlement/testing/FakeURLStreamHandler.java
@@ -18,6 +18,7 @@ package com.android.libraries.entitlement.testing;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableMap;
+
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -31,6 +32,7 @@ import java.security.cert.Certificate;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+
 import javax.annotation.Nullable;
 import javax.net.ssl.HttpsURLConnection;
 
@@ -40,137 +42,150 @@ import javax.net.ssl.HttpsURLConnection;
  */
 public class FakeURLStreamHandler extends URLStreamHandler implements URLStreamHandlerFactory {
 
-  private Map<String, FakeResponse> response;
-
-  private static final String ACCESS_TOKEN = "8dGozfI6%2FEaSsE7LaTfJKwdy";
-  private static final String LOCATION = "Location";
-  private static final String CONTENT_TYPE = "Content-Type";
-
-  @Override
-  public URLConnection openConnection(URL u) {
-    FakeHttpsURLConnection connection = new FakeHttpsURLConnection(u);
-    return connection;
-  }
-
-  @Override
-  public URLStreamHandler createURLStreamHandler(String protocol) {
-    return this;
-  }
-
-  public FakeURLStreamHandler stubResponse(Map<String, FakeResponse> response) {
-    this.response = response;
-    return this;
-  }
-
-  /** Fakes {@linkplain java.net.HttpURLConnection} to avoid making any network connection. */
-  public class FakeHttpsURLConnection extends HttpsURLConnection {
-
-    public ByteArrayOutputStream outputStream;
-
-    private final String urlString;
-
-    protected FakeHttpsURLConnection(URL url) {
-      super(url);
-      this.urlString = url.toString();
-    }
-
-    @Override
-    public InputStream getInputStream() throws IOException {
-      InputStream inputStream = new ByteArrayInputStream(response.get(urlString).responseBody());
-      if (inputStream == null) {
-        throw new IOException();
-      }
-      return inputStream;
-    }
-
-    @Override
-    public OutputStream getOutputStream() {
-      outputStream = new ByteArrayOutputStream();
-      return outputStream;
-    }
+    private Map<String, FakeResponse> mResponseMap;
 
-    @Override
-    public int getResponseCode() {
-      return response.get(urlString).responseCode();
-    }
+    private static final String ACCESS_TOKEN = "8dGozfI6%2FEaSsE7LaTfJKwdy";
+    private static final String LOCATION = "Location";
+    private static final String CONTENT_TYPE = "Content-Type";
+    private static final String RETRY_AFTER = "Retry-After";
 
     @Override
-    public Map<String, List<String>> getHeaderFields() {
-      List<String> locationList = new ArrayList<>();
-      locationList.add("access_token=" + ACCESS_TOKEN);
-      return ImmutableMap.of("Location", locationList);
+    public URLConnection openConnection(URL u) {
+        FakeHttpsURLConnection connection = new FakeHttpsURLConnection(u);
+        return connection;
     }
 
     @Override
-    public String getHeaderField(String name) {
-      switch (name) {
-        case LOCATION:
-          return "Location: " + response.get(urlString).responseLocation();
-        case CONTENT_TYPE:
-          return response.get(urlString).contentType();
-        default:
-          return null;
-      }
+    public URLStreamHandler createURLStreamHandler(String protocol) {
+        return this;
     }
 
-    @Override
-    public void connect() {}
-
-    @Override
-    public void disconnect() {}
-
-    @Override
-    public boolean usingProxy() {
-      return false;
+    public FakeURLStreamHandler stubResponse(Map<String, FakeResponse> response) {
+        this.mResponseMap = response;
+        return this;
     }
 
-    @Override
-    public String getCipherSuite() {
-      return null;
+    /**
+     * Fakes {@linkplain java.net.HttpURLConnection} to avoid making any network connection.
+     */
+    public class FakeHttpsURLConnection extends HttpsURLConnection {
+
+        public ByteArrayOutputStream mByteArrayOutputStream;
+
+        private final String mUrlString;
+
+        protected FakeHttpsURLConnection(URL url) {
+            super(url);
+            this.mUrlString = url.toString();
+        }
+
+        @Override
+        public InputStream getInputStream() throws IOException {
+            InputStream inputStream = new ByteArrayInputStream(
+                    mResponseMap.get(mUrlString).responseBody());
+            if (inputStream == null) {
+                throw new IOException();
+            }
+            return inputStream;
+        }
+
+        @Override
+        public OutputStream getOutputStream() {
+            mByteArrayOutputStream = new ByteArrayOutputStream();
+            return mByteArrayOutputStream;
+        }
+
+        @Override
+        public int getResponseCode() {
+            return mResponseMap.get(mUrlString).responseCode();
+        }
+
+        @Override
+        public Map<String, List<String>> getHeaderFields() {
+            List<String> locationList = new ArrayList<>();
+            locationList.add("access_token=" + ACCESS_TOKEN);
+            return ImmutableMap.of("Location", locationList);
+        }
+
+        @Override
+        public String getHeaderField(String name) {
+            switch (name) {
+                case LOCATION:
+                    return "Location: " + mResponseMap.get(mUrlString).responseLocation();
+                case CONTENT_TYPE:
+                    return mResponseMap.get(mUrlString).contentType();
+                case RETRY_AFTER:
+                    return mResponseMap.get(mUrlString).retryAfter();
+                default:
+                    return null;
+            }
+        }
+
+        @Override
+        public void connect() {
+        }
+
+        @Override
+        public void disconnect() {
+        }
+
+        @Override
+        public boolean usingProxy() {
+            return false;
+        }
+
+        @Override
+        public String getCipherSuite() {
+            return null;
+        }
+
+        @Override
+        public Certificate[] getLocalCertificates() {
+            return null;
+        }
+
+        @Override
+        public Certificate[] getServerCertificates() {
+            return null;
+        }
     }
 
-    @Override
-    public Certificate[] getLocalCertificates() {
-      return null;
-    }
+    @AutoValue
+    public abstract static class FakeResponse {
+        public abstract int responseCode();
 
-    @Override
-    public Certificate[] getServerCertificates() {
-      return null;
-    }
-  }
+        @Nullable
+        public abstract String responseLocation();
 
-  @AutoValue
-  public abstract static class FakeResponse {
-    public abstract int responseCode();
+        @SuppressWarnings("mutable") // For test only
+        public abstract byte[] responseBody();
 
-    @Nullable
-    public abstract String responseLocation();
+        public abstract String contentType();
 
-    @SuppressWarnings("mutable") // For test only
-    public abstract byte[] responseBody();
+        public abstract String retryAfter();
 
-    public abstract String contentType();
+        public static Builder builder() {
+            return new AutoValue_FakeURLStreamHandler_FakeResponse.Builder()
+                    .setResponseBody(new byte[]{})
+                    .setContentType("")
+                    .setResponseCode(0)
+                    .setResponseLocation("")
+                    .setRetryAfter("");
+        }
 
-    public static Builder builder() {
-      return new AutoValue_FakeURLStreamHandler_FakeResponse.Builder()
-          .setResponseBody(new byte[] {})
-          .setContentType("")
-          .setResponseCode(0)
-          .setResponseLocation("");
-    }
+        @AutoValue.Builder
+        public abstract static class Builder {
+            public abstract Builder setResponseCode(int value);
 
-    @AutoValue.Builder
-    public abstract static class Builder {
-      public abstract Builder setResponseCode(int value);
+            public abstract Builder setResponseLocation(String value);
 
-      public abstract Builder setResponseLocation(String value);
+            public abstract Builder setResponseBody(byte[] value);
 
-      public abstract Builder setResponseBody(byte[] value);
+            public abstract Builder setContentType(String contentType);
 
-      public abstract Builder setContentType(String contentType);
+            public abstract Builder setRetryAfter(String retryAfter);
 
-      public abstract FakeResponse build();
+            public abstract FakeResponse build();
+        }
     }
-  }
 }
\ No newline at end of file