Add public getModelDownloadId method (#2315)

* Add public getModelDownloadId method - can be used to track download progress.

* Apply suggestions from code review

Co-authored-by: Rodrigo Lazo <[email protected]>

Co-authored-by: Rodrigo Lazo <[email protected]>

Author: 2 people

firebase-ml-modeldownloader/api.txt

@@ -57,6 +57,7 @@ package com.google.firebase.ml.modeldownloader {
     method @NonNull public static com.google.firebase.ml.modeldownloader.FirebaseModelDownloader getInstance();
     method @NonNull public static com.google.firebase.ml.modeldownloader.FirebaseModelDownloader getInstance(@NonNull com.google.firebase.FirebaseApp);
     method @NonNull public com.google.android.gms.tasks.Task<com.google.firebase.ml.modeldownloader.CustomModel> getModel(@NonNull String, @NonNull com.google.firebase.ml.modeldownloader.DownloadType, @Nullable com.google.firebase.ml.modeldownloader.CustomModelDownloadConditions);
+    method public long getModelDownloadId(@NonNull String);
     method @NonNull public com.google.android.gms.tasks.Task<java.util.Set<com.google.firebase.ml.modeldownloader.CustomModel>> listDownloadedModels();
     method public void setStatsCollectionEnabled(boolean);
   }

firebase-ml-modeldownloader/src/main/java/com/google/firebase/ml/modeldownloader/CustomModelDownloadConditions.java

@@ -54,7 +54,7 @@ public static class Builder {
     private boolean isWifiRequired = false;
     private boolean isDeviceIdleRequired = false;
 
-    /** Sets whether charging is required. Only works on Android N and above. */
+    /** Sets charging as required. Only works on Android N and above. */
     @NonNull
     @RequiresApi(VERSION_CODES.N)
     @TargetApi(VERSION_CODES.N)
@@ -63,15 +63,15 @@ public Builder requireCharging() {
       return this;
     }
 
-    /** Sets whether wifi is required. */
+    /** Sets wifi as required. */
     @NonNull
     public Builder requireWifi() {
       this.isWifiRequired = true;
       return this;
     }
 
     /**
-     * Sets whether device idle is required.
+     * Sets device idle as required.
      *
      * <p>Idle mode is a loose definition provided by the system, which means that the device is not
      * in use, and has not been in use for some time.

firebase-ml-modeldownloader/src/main/java/com/google/firebase/ml/modeldownloader/FirebaseModelDownloader.java

@@ -422,6 +422,31 @@ public void setStatsCollectionEnabled(boolean enabled) {
     sharedPreferencesUtil.setCustomModelStatsCollectionEnabled(enabled);
   }
 
+  /**
+   * Get the current models' download id (returns background download id when applicable). This id
+   * can be used to create a progress bar to track file download progress.
+   *
+   * <p>If no model exists or there is no download in progress, return 0.
+   *
+   * <p>If 0 is returned immediately after starting a download via getModel, then
+   *
+   * <ul>
+   *   <li>the enqueuing wasn't needed: the getModel task already completed and/or no background
+   *       update.
+   *   <li>the enqueuing hasn't completed: the download id hasn't generated yet - try again.
+   * </ul>
+   *
+   * @param modelName - model name
+   * @return id associated with Android Download Manager.
+   */
+  public long getModelDownloadId(@NonNull String modelName) {
+    CustomModel localModel = sharedPreferencesUtil.getDownloadingCustomModelDetails(modelName);
+    if (localModel != null) {
+      return localModel.getDownloadId();
+    }
+    return 0;
+  }
+
   /** Returns the nick name of the {@link FirebaseApp} of this {@link FirebaseModelDownloader} */
   @VisibleForTesting
   String getApplicationId() {

firebase-ml-modeldownloader/src/test/java/com/google/firebase/ml/modeldownloader/FirebaseModelDownloaderTest.java

@@ -725,4 +725,23 @@ public void setStatsCollectionEnabled() {
     firebaseModelDownloader.setStatsCollectionEnabled(false);
     verify(mockPrefs, times(1)).setCustomModelStatsCollectionEnabled(eq(false));
   }
+
+  @Test
+  public void getModelDownloadId_noDownload() {
+    when(mockPrefs.getDownloadingCustomModelDetails(eq(MODEL_NAME))).thenReturn(customModelLoaded);
+    assertEquals(firebaseModelDownloader.getModelDownloadId(MODEL_NAME), 0);
+  }
+
+  @Test
+  public void getModelDownloadId_noNamedModel() {
+    when(mockPrefs.getDownloadingCustomModelDetails(eq(MODEL_NAME))).thenReturn(null);
+    assertEquals(firebaseModelDownloader.getModelDownloadId(MODEL_NAME), 0);
+  }
+
+  @Test
+  public void getModelDownloadId_download() {
+    when(mockPrefs.getDownloadingCustomModelDetails(eq(MODEL_NAME)))
+        .thenReturn(UPDATE_IN_PROGRESS_CUSTOM_MODEL);
+    assertEquals(firebaseModelDownloader.getModelDownloadId(MODEL_NAME), DOWNLOAD_ID);
+  }
 }