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.common.statsd; 18 19 import android.util.Log; 20 import androidx.test.core.app.ApplicationProvider; 21 import com.android.internal.os.StatsdConfigProto.StatsdConfig; 22 import com.android.os.AtomsProto.Atom; 23 import com.android.os.AtomsProto.TextClassifierDownloadReported; 24 import com.android.os.AtomsProto.TextClassifierDownloadWorkCompleted; 25 import com.android.os.AtomsProto.TextClassifierDownloadWorkScheduled; 26 import com.google.common.collect.ImmutableList; 27 import java.util.stream.Collectors; 28 import org.junit.rules.ExternalResource; 29 30 // TODO(licha): Make this generic and useful for other atoms. 31 /** Test rule to set up/clean up statsd for download logger tests. */ 32 public final class TextClassifierDownloadLoggerTestRule extends ExternalResource { 33 private static final String TAG = "DownloadLoggerTestRule"; 34 35 // Statsd config IDs, which are arbitrary. 36 private static final long CONFIG_ID_DOWNLOAD_REPORTED = 423779; 37 private static final long CONFIG_ID_DOWNLOAD_WORK_SCHEDULED = 42; 38 private static final long CONFIG_ID_DOWNLOAD_WORK_COMPLETED = 2021; 39 40 private static final long SHORT_TIMEOUT_MS = 1000; 41 42 @Override before()43 public void before() throws Exception { 44 StatsdTestUtils.cleanup(CONFIG_ID_DOWNLOAD_REPORTED); 45 StatsdTestUtils.cleanup(CONFIG_ID_DOWNLOAD_WORK_SCHEDULED); 46 StatsdTestUtils.cleanup(CONFIG_ID_DOWNLOAD_WORK_COMPLETED); 47 48 StatsdConfig.Builder builder1 = 49 StatsdConfig.newBuilder() 50 .setId(CONFIG_ID_DOWNLOAD_REPORTED) 51 .addAllowedLogSource(ApplicationProvider.getApplicationContext().getPackageName()); 52 StatsdTestUtils.addAtomMatcher(builder1, Atom.TEXT_CLASSIFIER_DOWNLOAD_REPORTED_FIELD_NUMBER); 53 StatsdTestUtils.pushConfig(builder1.build()); 54 55 StatsdConfig.Builder builder2 = 56 StatsdConfig.newBuilder() 57 .setId(CONFIG_ID_DOWNLOAD_WORK_SCHEDULED) 58 .addAllowedLogSource(ApplicationProvider.getApplicationContext().getPackageName()); 59 StatsdTestUtils.addAtomMatcher( 60 builder2, Atom.TEXT_CLASSIFIER_DOWNLOAD_WORK_SCHEDULED_FIELD_NUMBER); 61 StatsdTestUtils.pushConfig(builder2.build()); 62 63 StatsdConfig.Builder builder3 = 64 StatsdConfig.newBuilder() 65 .setId(CONFIG_ID_DOWNLOAD_WORK_COMPLETED) 66 .addAllowedLogSource(ApplicationProvider.getApplicationContext().getPackageName()); 67 StatsdTestUtils.addAtomMatcher( 68 builder3, Atom.TEXT_CLASSIFIER_DOWNLOAD_WORK_COMPLETED_FIELD_NUMBER); 69 StatsdTestUtils.pushConfig(builder3.build()); 70 } 71 72 @Override after()73 public void after() { 74 try { 75 StatsdTestUtils.cleanup(CONFIG_ID_DOWNLOAD_REPORTED); 76 StatsdTestUtils.cleanup(CONFIG_ID_DOWNLOAD_WORK_SCHEDULED); 77 StatsdTestUtils.cleanup(CONFIG_ID_DOWNLOAD_WORK_COMPLETED); 78 } catch (Exception e) { 79 Log.e(TAG, "Failed to clean up statsd after tests."); 80 } 81 } 82 83 /** 84 * Gets a list of TextClassifierDownloadReported atoms written into statsd, sorted by increasing 85 * timestamp. 86 */ getLoggedDownloadReportedAtoms()87 public ImmutableList<TextClassifierDownloadReported> getLoggedDownloadReportedAtoms() 88 throws Exception { 89 ImmutableList<Atom> loggedAtoms = 90 StatsdTestUtils.getLoggedAtoms(CONFIG_ID_DOWNLOAD_REPORTED, SHORT_TIMEOUT_MS); 91 return ImmutableList.copyOf( 92 loggedAtoms.stream() 93 .filter(Atom::hasTextClassifierDownloadReported) 94 .map(Atom::getTextClassifierDownloadReported) 95 .collect(Collectors.toList())); 96 } 97 98 /** 99 * Gets a list of TextClassifierDownloadWorkScheduled atoms written into statsd, sorted by 100 * increasing timestamp. 101 */ getLoggedDownloadWorkScheduledAtoms()102 public ImmutableList<TextClassifierDownloadWorkScheduled> getLoggedDownloadWorkScheduledAtoms() 103 throws Exception { 104 ImmutableList<Atom> loggedAtoms = 105 StatsdTestUtils.getLoggedAtoms(CONFIG_ID_DOWNLOAD_WORK_SCHEDULED, SHORT_TIMEOUT_MS); 106 return ImmutableList.copyOf( 107 loggedAtoms.stream() 108 .filter(Atom::hasTextClassifierDownloadWorkScheduled) 109 .map(Atom::getTextClassifierDownloadWorkScheduled) 110 .collect(Collectors.toList())); 111 } 112 113 /** 114 * Gets a list of TextClassifierDownloadWorkCompleted atoms written into statsd, sorted by 115 * increasing timestamp. 116 */ getLoggedDownloadWorkCompletedAtoms()117 public ImmutableList<TextClassifierDownloadWorkCompleted> getLoggedDownloadWorkCompletedAtoms() 118 throws Exception { 119 ImmutableList<Atom> loggedAtoms = 120 StatsdTestUtils.getLoggedAtoms(CONFIG_ID_DOWNLOAD_WORK_COMPLETED, SHORT_TIMEOUT_MS); 121 return ImmutableList.copyOf( 122 loggedAtoms.stream() 123 .filter(Atom::hasTextClassifierDownloadWorkCompleted) 124 .map(Atom::getTextClassifierDownloadWorkCompleted) 125 .collect(Collectors.toList())); 126 } 127 } 128