Add tests for ServiceEntitment

Bug: 173450048
Test: presubmit
Change-Id: Iaf9435e3c3c42295c2cc33a33fa5e4a7af20dcba

5 files changed, 209 insertions, 7 deletions

diff --git a/java/com/android/libraries/entitlement/ServiceEntitlement.java b/java/com/android/libraries/entitlement/ServiceEntitlement.java
index c196783..abf6ace 100644
--- a/java/com/android/libraries/entitlement/ServiceEntitlement.java
+++ b/java/com/android/libraries/entitlement/ServiceEntitlement.java
@@ -20,6 +20,8 @@ import android.content.Context;
 
 import com.android.libraries.entitlement.eapaka.EapAkaApi;
 
+import com.google.common.annotations.VisibleForTesting;
+
 import java.util.List;
 
 import androidx.annotation.Nullable;
@@ -41,8 +43,8 @@ public class ServiceEntitlement {
     public static final String APP_ODSA_PRIMARY = "ap2009";
 
     private final Context context;
-    private final int simSubscriptionId;
     private final CarrierConfig carrierConfig;
+    private final EapAkaApi eapAkaApi;
 
     /**
      * Creates an instance for service entitlement configuration query and operation for the
@@ -58,8 +60,15 @@ public class ServiceEntitlement {
      */
     public ServiceEntitlement(Context context, CarrierConfig carrierConfig, int simSubscriptionId) {
         this.context = context;
-        this.simSubscriptionId = simSubscriptionId;
         this.carrierConfig = carrierConfig;
+        this.eapAkaApi = new EapAkaApi(context, simSubscriptionId);
+    }
+
+    @VisibleForTesting
+    ServiceEntitlement(Context context, CarrierConfig carrierConfig, EapAkaApi eapAkaApi) {
+        this.context = context;
+        this.carrierConfig = carrierConfig;
+        this.eapAkaApi = eapAkaApi;
     }
 
     /**
@@ -105,7 +114,6 @@ public class ServiceEntitlement {
     @Nullable
     public String queryEntitlementStatus(String appId, ServiceEntitlementRequest request)
             throws ServiceEntitlementException {
-        EapAkaApi eapAkaApi = new EapAkaApi(context, simSubscriptionId);
         return eapAkaApi.queryEntitlementStatus(appId, carrierConfig.serverUrl(), request);
     }
 
diff --git a/java/com/android/libraries/entitlement/eapaka/EapAkaApi.java b/java/com/android/libraries/entitlement/eapaka/EapAkaApi.java
index 898b782..b9eaadd 100644
--- a/java/com/android/libraries/entitlement/eapaka/EapAkaApi.java
+++ b/java/com/android/libraries/entitlement/eapaka/EapAkaApi.java
@@ -35,6 +35,7 @@ import org.json.JSONObject;
 
 import androidx.annotation.Nullable;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.net.HttpHeaders;
 
 import java.net.CookieHandler;
@@ -86,6 +87,13 @@ public class EapAkaApi {
         this.httpClient = new HttpClient();
     }
 
+    @VisibleForTesting
+    EapAkaApi(Context context, int simSubscriptionId, HttpClient httpClient) {
+        this.context = context;
+        this.simSubscriptionId = simSubscriptionId;
+        this.httpClient = httpClient;
+    }
+
     /**
      * Retrieves raw entitlement configuration doc though EAP-AKA authentication.
      *
@@ -160,11 +168,12 @@ public class EapAkaApi {
         return response.body();
     }
 
-    private String entitlementStatusUrl(
+    @VisibleForTesting
+    String entitlementStatusUrl(
             String appId, String serverUrl, ServiceEntitlementRequest request) {
         TelephonyManager telephonyManager =
-                context.getSystemService(TelephonyManager.class).createForSubscriptionId(
-                        simSubscriptionId);
+                (TelephonyManager) context.getSystemService(Context.TELEPHONY_SERVICE);
+        telephonyManager = telephonyManager.createForSubscriptionId(simSubscriptionId);
         Uri.Builder urlBuilder = Uri.parse(serverUrl).buildUpon();
         if (TextUtils.isEmpty(request.authenticationToken())) {
             // EAP_ID required for initial AuthN
diff --git a/java/com/android/libraries/entitlement/eapaka/EapAkaResponse.java b/java/com/android/libraries/entitlement/eapaka/EapAkaResponse.java
index 5021317..8f4f138 100644
--- a/java/com/android/libraries/entitlement/eapaka/EapAkaResponse.java
+++ b/java/com/android/libraries/entitlement/eapaka/EapAkaResponse.java
@@ -25,6 +25,8 @@ import android.util.Log;
 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;
 
@@ -73,6 +75,9 @@ class EapAkaResponse {
 
     private boolean valid;
 
+    @VisibleForTesting
+    static String challengeResponseForTesting;
+
     public EapAkaResponse(String eapAkaChallenge) {
         try {
             parseEapAkaChallengeRequest(eapAkaChallenge);
@@ -201,6 +206,10 @@ class EapAkaResponse {
      */
     public String getEapAkaChallengeResponse(Context context, int simSubscriptionId)
             throws ServiceEntitlementException {
+        if(challengeResponseForTesting != null) {
+            return challengeResponseForTesting;
+        }
+
         if (!valid) {
             throw new ServiceEntitlementException("EAP-AKA Challenge message not valid!");
         }
diff --git a/java/com/android/libraries/entitlement/eapaka/MasterKey.java b/java/com/android/libraries/entitlement/eapaka/MasterKey.java
index 652fa65..57b8ab3 100644
--- a/java/com/android/libraries/entitlement/eapaka/MasterKey.java
+++ b/java/com/android/libraries/entitlement/eapaka/MasterKey.java
@@ -93,7 +93,7 @@ class MasterKey {
             messageDigest.update(data);
             masterKey = messageDigest.digest();
         } catch (NoSuchAlgorithmException e) {
-            Log.d(TAG, "process sHA-1 failed", e);
+            Log.d(TAG, "process SHA-1 failed", e);
         }
 
         // Generate TEKs
diff --git a/java/com/android/libraries/entitlement/testing/FakeURLStreamHandler.java b/java/com/android/libraries/entitlement/testing/FakeURLStreamHandler.java
new file mode 100644
index 0000000..0763747
--- /dev/null
+++ b/java/com/android/libraries/entitlement/testing/FakeURLStreamHandler.java
@@ -0,0 +1,176 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.net.URL;
+import java.net.URLConnection;
+import java.net.URLStreamHandler;
+import java.net.URLStreamHandlerFactory;
+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;
+
+/**
+ * Fakes {@linkplain java.net.URLStreamHandler} which is used to set URLStreamHandlerFactory for URL
+ * as {@linkplain java.net.URL} is a final class and cannot be mocked using mockito.
+ */
+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;
+    }
+
+    @Override
+    public int getResponseCode() {
+      return response.get(urlString).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: " + response.get(urlString).responseLocation();
+        case CONTENT_TYPE:
+          return response.get(urlString).contentType();
+        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;
+    }
+  }
+
+  @AutoValue
+  public abstract static class FakeResponse {
+    public abstract int responseCode();
+
+    @Nullable
+    public abstract String responseLocation();
+
+    @SuppressWarnings("mutable") // For test only
+    public abstract byte[] responseBody();
+
+    public abstract String contentType();
+
+    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);
+
+      public abstract Builder setResponseLocation(String value);
+
+      public abstract Builder setResponseBody(byte[] value);
+
+      public abstract Builder setContentType(String contentType);
+
+      public abstract FakeResponse build();
+    }
+  }
+}
\ No newline at end of file