1 /*
2  * Copyright (C) 2018 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.textclassifier.downloader;
18 
19 import static com.google.common.truth.Truth.assertThat;
20 import static org.testng.Assert.expectThrows;
21 
22 import android.content.Context;
23 import androidx.room.Room;
24 import androidx.test.core.app.ApplicationProvider;
25 import androidx.test.ext.junit.runners.AndroidJUnit4;
26 import com.android.textclassifier.common.ModelType;
27 import com.android.textclassifier.downloader.DownloadedModelDatabase.Manifest;
28 import com.android.textclassifier.downloader.DownloadedModelDatabase.ManifestEnrollment;
29 import com.android.textclassifier.downloader.DownloadedModelDatabase.ManifestModelCrossRef;
30 import com.android.textclassifier.downloader.DownloadedModelDatabase.Model;
31 import com.android.textclassifier.downloader.DownloadedModelDatabase.ModelView;
32 import com.google.common.collect.ImmutableList;
33 import com.google.common.collect.Iterables;
34 import java.io.IOException;
35 import java.util.List;
36 import org.junit.After;
37 import org.junit.Before;
38 import org.junit.Test;
39 import org.junit.runner.RunWith;
40 
41 @RunWith(AndroidJUnit4.class)
42 public class DownloadedModelDatabaseTest {
43   private static final String MODEL_URL = "https://model.url";
44   private static final String MODEL_URL_2 = "https://model2.url";
45   private static final String MODEL_PATH = "/data/test.model";
46   private static final String MODEL_PATH_2 = "/data/test.model2";
47   private static final String MANIFEST_URL = "https://manifest.url";
48   private static final String MANIFEST_URL_2 = "https://manifest2.url";
49   private static final String MODEL_TYPE = ModelType.ANNOTATOR;
50   private static final String MODEL_TYPE_2 = ModelType.ACTIONS_SUGGESTIONS;
51   private static final String LOCALE_TAG = "zh";
52 
53   private DownloadedModelDatabase db;
54 
55   @Before
createDb()56   public void createDb() {
57     Context context = ApplicationProvider.getApplicationContext();
58     db = Room.inMemoryDatabaseBuilder(context, DownloadedModelDatabase.class).build();
59   }
60 
61   @After
closeDb()62   public void closeDb() throws IOException {
63     db.close();
64   }
65 
66   @Test
insertModelAndRead()67   public void insertModelAndRead() throws Exception {
68     Model model = Model.create(MODEL_URL, MODEL_PATH);
69     db.dao().insert(model);
70     List<Model> models = db.dao().queryAllModels();
71     assertThat(models).containsExactly(model);
72   }
73 
74   @Test
insertModelAndDelete()75   public void insertModelAndDelete() throws Exception {
76     Model model = Model.create(MODEL_URL, MODEL_PATH);
77     db.dao().insert(model);
78     db.dao().deleteModels(ImmutableList.of(model));
79     List<Model> models = db.dao().queryAllModels();
80     assertThat(models).isEmpty();
81   }
82 
83   @Test
insertManifestAndRead()84   public void insertManifestAndRead() throws Exception {
85     Manifest manifest =
86         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
87     db.dao().insert(manifest);
88     List<Manifest> manifests = db.dao().queryAllManifests();
89     assertThat(manifests).containsExactly(manifest);
90   }
91 
92   @Test
insertManifestAndDelete()93   public void insertManifestAndDelete() throws Exception {
94     Manifest manifest =
95         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
96     db.dao().insert(manifest);
97     db.dao().deleteManifests(ImmutableList.of(manifest));
98     List<Manifest> manifests = db.dao().queryAllManifests();
99     assertThat(manifests).isEmpty();
100   }
101 
102   @Test
insertManifestModelCrossRefAndRead()103   public void insertManifestModelCrossRefAndRead() throws Exception {
104     Model model = Model.create(MODEL_URL, MODEL_PATH);
105     db.dao().insert(model);
106     Manifest manifest =
107         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
108     db.dao().insert(manifest);
109     ManifestModelCrossRef manifestModelCrossRef =
110         ManifestModelCrossRef.create(MANIFEST_URL, MODEL_URL);
111     db.dao().insert(manifestModelCrossRef);
112     List<ManifestModelCrossRef> manifestModelCrossRefs = db.dao().queryAllManifestModelCrossRefs();
113     assertThat(manifestModelCrossRefs).containsExactly(manifestModelCrossRef);
114   }
115 
116   @Test
insertManifestModelCrossRefAndDelete()117   public void insertManifestModelCrossRefAndDelete() throws Exception {
118     Model model = Model.create(MODEL_URL, MODEL_PATH);
119     db.dao().insert(model);
120     Manifest manifest =
121         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
122     db.dao().insert(manifest);
123     ManifestModelCrossRef manifestModelCrossRef =
124         ManifestModelCrossRef.create(MANIFEST_URL, MODEL_URL);
125     db.dao().insert(manifestModelCrossRef);
126     db.dao().deleteManifestModelCrossRefs(ImmutableList.of(manifestModelCrossRef));
127     List<ManifestModelCrossRef> manifestModelCrossRefs = db.dao().queryAllManifestModelCrossRefs();
128     assertThat(manifestModelCrossRefs).isEmpty();
129   }
130 
131   @Test
insertManifestModelCrossRefAndDeleteManifest()132   public void insertManifestModelCrossRefAndDeleteManifest() throws Exception {
133     Model model = Model.create(MODEL_URL, MODEL_PATH);
134     db.dao().insert(model);
135     Manifest manifest =
136         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
137     db.dao().insert(manifest);
138     ManifestModelCrossRef manifestModelCrossRef =
139         ManifestModelCrossRef.create(MANIFEST_URL, MODEL_URL);
140     db.dao().insert(manifestModelCrossRef);
141     db.dao().deleteManifests(ImmutableList.of(manifest)); // ON CASCADE
142     List<ManifestModelCrossRef> manifestModelCrossRefs = db.dao().queryAllManifestModelCrossRefs();
143     assertThat(manifestModelCrossRefs).isEmpty();
144   }
145 
146   @Test
insertManifestModelCrossRefAndDeleteModel()147   public void insertManifestModelCrossRefAndDeleteModel() throws Exception {
148     Model model = Model.create(MODEL_URL, MODEL_PATH);
149     db.dao().insert(model);
150     Manifest manifest =
151         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
152     db.dao().insert(manifest);
153     ManifestModelCrossRef manifestModelCrossRef =
154         ManifestModelCrossRef.create(MANIFEST_URL, MODEL_URL);
155     db.dao().insert(manifestModelCrossRef);
156     db.dao().deleteModels(ImmutableList.of(model)); // ON CASCADE
157     List<ManifestModelCrossRef> manifestModelCrossRefs = db.dao().queryAllManifestModelCrossRefs();
158     assertThat(manifestModelCrossRefs).isEmpty();
159   }
160 
161   @Test
insertManifestModelCrossRefWithoutManifest()162   public void insertManifestModelCrossRefWithoutManifest() throws Exception {
163     Model model = Model.create(MODEL_URL, MODEL_PATH);
164     db.dao().insert(model);
165     ManifestModelCrossRef manifestModelCrossRef =
166         ManifestModelCrossRef.create(MANIFEST_URL, MODEL_URL);
167     expectThrows(Throwable.class, () -> db.dao().insert(manifestModelCrossRef));
168   }
169 
170   @Test
insertManifestModelCrossRefWithoutModel()171   public void insertManifestModelCrossRefWithoutModel() throws Exception {
172     Manifest manifest =
173         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
174     db.dao().insert(manifest);
175     ManifestModelCrossRef manifestModelCrossRef =
176         ManifestModelCrossRef.create(MANIFEST_URL, MODEL_URL);
177     expectThrows(Throwable.class, () -> db.dao().insert(manifestModelCrossRef));
178   }
179 
180   @Test
insertManifestEnrollmentAndRead()181   public void insertManifestEnrollmentAndRead() throws Exception {
182     Manifest manifest =
183         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
184     db.dao().insert(manifest);
185     ManifestEnrollment manifestEnrollment =
186         ManifestEnrollment.create(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL);
187     db.dao().insert(manifestEnrollment);
188     List<ManifestEnrollment> manifestEnrollments = db.dao().queryAllManifestEnrollments();
189     assertThat(manifestEnrollments).containsExactly(manifestEnrollment);
190   }
191 
192   @Test
insertManifestEnrollmentAndDelete()193   public void insertManifestEnrollmentAndDelete() throws Exception {
194     Manifest manifest =
195         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
196     db.dao().insert(manifest);
197     ManifestEnrollment manifestEnrollment =
198         ManifestEnrollment.create(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL);
199     db.dao().insert(manifestEnrollment);
200     db.dao().deleteManifestEnrollments(ImmutableList.of(manifestEnrollment));
201     List<ManifestEnrollment> manifestEnrollments = db.dao().queryAllManifestEnrollments();
202     assertThat(manifestEnrollments).isEmpty();
203   }
204 
205   @Test
insertManifestEnrollmentAndDeleteManifest()206   public void insertManifestEnrollmentAndDeleteManifest() throws Exception {
207     Manifest manifest =
208         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
209     db.dao().insert(manifest);
210     ManifestEnrollment manifestEnrollment =
211         ManifestEnrollment.create(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL);
212     db.dao().insert(manifestEnrollment);
213     db.dao().deleteManifests(ImmutableList.of(manifest));
214     List<ManifestEnrollment> manifestEnrollments = db.dao().queryAllManifestEnrollments();
215     assertThat(manifestEnrollments).isEmpty();
216   }
217 
218   @Test
insertManifestEnrollmentWithoutManifest()219   public void insertManifestEnrollmentWithoutManifest() throws Exception {
220     ManifestEnrollment manifestEnrollment =
221         ManifestEnrollment.create(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL);
222     expectThrows(Throwable.class, () -> db.dao().insert(manifestEnrollment));
223   }
224 
225   @Test
insertModelViewAndRead()226   public void insertModelViewAndRead() throws Exception {
227     Model model = Model.create(MODEL_URL, MODEL_PATH);
228     db.dao().insert(model);
229     Manifest manifest =
230         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
231     db.dao().insert(manifest);
232     ManifestModelCrossRef manifestModelCrossRef =
233         ManifestModelCrossRef.create(MANIFEST_URL, MODEL_URL);
234     db.dao().insert(manifestModelCrossRef);
235     ManifestEnrollment manifestEnrollment =
236         ManifestEnrollment.create(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL);
237     db.dao().insert(manifestEnrollment);
238 
239     List<ModelView> modelViews = db.dao().queryAllModelViews();
240     ModelView modelView = Iterables.getOnlyElement(modelViews);
241     assertThat(modelView.getManifestEnrollment()).isEqualTo(manifestEnrollment);
242     assertThat(modelView.getModel()).isEqualTo(model);
243   }
244 
245   @Test
queryModelWithModelUrl()246   public void queryModelWithModelUrl() throws Exception {
247     Model model = Model.create(MODEL_URL, MODEL_PATH);
248     db.dao().insert(model);
249     Model model2 = Model.create(MODEL_URL_2, MODEL_PATH_2);
250     db.dao().insert(model2);
251 
252     assertThat(db.dao().queryModelWithModelUrl(MODEL_URL)).containsExactly(model);
253     assertThat(db.dao().queryModelWithModelUrl(MODEL_URL_2)).containsExactly(model2);
254   }
255 
256   @Test
queryManifestWithManifestUrl()257   public void queryManifestWithManifestUrl() throws Exception {
258     Manifest manifest =
259         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
260     db.dao().insert(manifest);
261     Manifest manifest2 =
262         Manifest.create(MANIFEST_URL_2, Manifest.STATUS_FAILED, /* failureCounts= */ 1);
263     db.dao().insert(manifest2);
264 
265     assertThat(db.dao().queryManifestWithManifestUrl(MANIFEST_URL)).containsExactly(manifest);
266     assertThat(db.dao().queryManifestWithManifestUrl(MANIFEST_URL_2)).containsExactly(manifest2);
267   }
268 
269   @Test
queryManifestEnrollmentWithModelTypeAndLocaleTag()270   public void queryManifestEnrollmentWithModelTypeAndLocaleTag() throws Exception {
271     Manifest manifest =
272         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
273     db.dao().insert(manifest);
274     Manifest manifest2 =
275         Manifest.create(MANIFEST_URL_2, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
276     db.dao().insert(manifest2);
277     ManifestEnrollment manifestEnrollment =
278         ManifestEnrollment.create(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL);
279     db.dao().insert(manifestEnrollment);
280     ManifestEnrollment manifestEnrollment2 =
281         ManifestEnrollment.create(MODEL_TYPE_2, LOCALE_TAG, MANIFEST_URL_2);
282     db.dao().insert(manifestEnrollment2);
283 
284     assertThat(db.dao().queryManifestEnrollmentWithModelTypeAndLocaleTag(MODEL_TYPE, LOCALE_TAG))
285         .containsExactly(manifestEnrollment);
286     assertThat(db.dao().queryManifestEnrollmentWithModelTypeAndLocaleTag(MODEL_TYPE_2, LOCALE_TAG))
287         .containsExactly(manifestEnrollment2);
288   }
289 
290   @Test
insertManifestAndModelCrossRef()291   public void insertManifestAndModelCrossRef() throws Exception {
292     Model model = Model.create(MODEL_URL, MODEL_PATH);
293     db.dao().insert(model);
294     Manifest manifest =
295         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
296     db.dao().insertManifestAndModelCrossRef(MANIFEST_URL, MODEL_URL);
297 
298     assertThat(db.dao().queryAllModels()).containsExactly(model);
299     assertThat(db.dao().queryAllManifests()).containsExactly(manifest);
300   }
301 
302   @Test
increaseManifestFailureCounts()303   public void increaseManifestFailureCounts() throws Exception {
304     db.dao().increaseManifestFailureCounts(MODEL_URL);
305     Manifest manifest = Iterables.getOnlyElement(db.dao().queryManifestWithManifestUrl(MODEL_URL));
306     assertThat(manifest.getStatus()).isEqualTo(Manifest.STATUS_FAILED);
307     assertThat(manifest.getFailureCounts()).isEqualTo(1);
308     db.dao().increaseManifestFailureCounts(MODEL_URL);
309     manifest = Iterables.getOnlyElement(db.dao().queryManifestWithManifestUrl(MODEL_URL));
310     assertThat(manifest.getStatus()).isEqualTo(Manifest.STATUS_FAILED);
311     assertThat(manifest.getFailureCounts()).isEqualTo(2);
312   }
313 
314   @Test
deleteUnusedManifestsAndModels_unusedManifestAndUnusedModel()315   public void deleteUnusedManifestsAndModels_unusedManifestAndUnusedModel() throws Exception {
316     Model model = Model.create(MODEL_URL, MODEL_PATH);
317     db.dao().insert(model);
318     Model model2 = Model.create(MODEL_URL_2, MODEL_PATH_2);
319     db.dao().insert(model2);
320     Manifest manifest =
321         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
322     db.dao().insert(manifest);
323     db.dao().insertManifestAndModelCrossRef(MANIFEST_URL, MODEL_URL);
324     Manifest manifest2 =
325         Manifest.create(MANIFEST_URL_2, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
326     db.dao().insert(manifest2);
327     db.dao().insertManifestAndModelCrossRef(MANIFEST_URL_2, MODEL_URL_2);
328     ManifestEnrollment manifestEnrollment =
329         ManifestEnrollment.create(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL);
330     db.dao().insert(manifestEnrollment);
331 
332     db.dao().deleteUnusedManifestsAndModels();
333     assertThat(db.dao().queryAllManifests()).containsExactly(manifest);
334     assertThat(db.dao().queryAllModels()).containsExactly(model);
335   }
336 
337   @Test
deleteUnusedManifestsAndModels_unusedManifestAndSharedModel()338   public void deleteUnusedManifestsAndModels_unusedManifestAndSharedModel() throws Exception {
339     Model model = Model.create(MODEL_URL, MODEL_PATH);
340     db.dao().insert(model);
341     Manifest manifest =
342         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
343     db.dao().insert(manifest);
344     db.dao().insertManifestAndModelCrossRef(MANIFEST_URL, MODEL_URL);
345     Manifest manifest2 =
346         Manifest.create(MANIFEST_URL_2, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
347     db.dao().insert(manifest2);
348     db.dao().insertManifestAndModelCrossRef(MANIFEST_URL_2, MODEL_URL);
349     ManifestEnrollment manifestEnrollment =
350         ManifestEnrollment.create(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL);
351     db.dao().insert(manifestEnrollment);
352 
353     db.dao().deleteUnusedManifestsAndModels();
354     assertThat(db.dao().queryAllManifests()).containsExactly(manifest);
355     assertThat(db.dao().queryAllModels()).containsExactly(model);
356   }
357 
358   @Test
deleteUnusedManifestsAndModels_failedManifest()359   public void deleteUnusedManifestsAndModels_failedManifest() throws Exception {
360     Manifest manifest =
361         Manifest.create(MANIFEST_URL, Manifest.STATUS_FAILED, /* failureCounts= */ 1);
362     db.dao().insert(manifest);
363 
364     db.dao().deleteUnusedManifestsAndModels();
365     assertThat(db.dao().queryAllManifests()).containsExactly(manifest);
366   }
367 
368   @Test
deleteUnusedManifestsAndModels_unusedModels()369   public void deleteUnusedManifestsAndModels_unusedModels() throws Exception {
370     Model model = Model.create(MODEL_URL, MODEL_PATH);
371     db.dao().insert(model);
372     Model model2 = Model.create(MODEL_URL_2, MODEL_PATH_2);
373     db.dao().insert(model2);
374     Manifest manifest =
375         Manifest.create(MANIFEST_URL, Manifest.STATUS_SUCCEEDED, /* failureCounts= */ 0);
376     db.dao().insert(manifest);
377     db.dao().insertManifestAndModelCrossRef(MANIFEST_URL, MODEL_URL);
378     ManifestEnrollment manifestEnrollment =
379         ManifestEnrollment.create(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL);
380     db.dao().insert(manifestEnrollment);
381 
382     db.dao().deleteUnusedManifestsAndModels();
383     assertThat(db.dao().queryAllModels()).containsExactly(model);
384   }
385 
386   @Test
deleteUnusedManifestFailureRecords()387   public void deleteUnusedManifestFailureRecords() throws Exception {
388     Manifest manifest =
389         Manifest.create(MANIFEST_URL, Manifest.STATUS_FAILED, /* failureCounts= */ 1);
390     db.dao().insert(manifest);
391     Manifest manifest2 =
392         Manifest.create(MANIFEST_URL_2, Manifest.STATUS_FAILED, /* failureCounts= */ 1);
393     db.dao().insert(manifest2);
394 
395     db.dao().deleteUnusedManifestFailureRecords(ImmutableList.of(MANIFEST_URL));
396     assertThat(db.dao().queryAllManifests()).containsExactly(manifest);
397   }
398 }
399