Merge "Tests for libstatspull"
diff --git a/apex/tests/libstatspull/Android.bp b/apex/tests/libstatspull/Android.bp
new file mode 100644
index 0000000..e813964
--- /dev/null
+++ b/apex/tests/libstatspull/Android.bp
@@ -0,0 +1,56 @@
+// Copyright (C) 2019 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.
+
+android_test {
+    name: "LibStatsPullTests",
+    static_libs: [
+        "androidx.test.rules",
+        "platformprotoslite",
+        "statsdprotolite",
+        "truth-prebuilt",
+    ],
+    libs: [
+        "android.test.runner.stubs",
+        "android.test.base.stubs",
+    ],
+    jni_libs: [
+        "libstatspull_testhelper",
+    ],
+    srcs: [
+        "src/**/*.java",
+        "protos/**/*.proto",
+        ],
+    test_suites: [
+        "general-tests",
+    ],
+    platform_apis: true,
+    privileged: true,
+    certificate: "platform",
+    compile_multilib: "both",
+}
+
+cc_library_shared {
+    name: "libstatspull_testhelper",
+    srcs: ["jni/stats_pull_helper.cpp"],
+    cflags: [
+        "-Wall",
+        "-Werror",
+    ],
+    shared_libs: [
+        "libbinder",
+        "libutils",
+        "libstatspull",
+        "libstatssocket",
+    ],
+}
\ No newline at end of file
diff --git a/apex/tests/libstatspull/AndroidManifest.xml b/apex/tests/libstatspull/AndroidManifest.xml
new file mode 100644
index 0000000..bffd400
--- /dev/null
+++ b/apex/tests/libstatspull/AndroidManifest.xml
@@ -0,0 +1,30 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+ * Copyright (C) 2020 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.
+ -->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+          package="com.android.internal.os.statsd.libstats" >
+
+
+    <uses-permission android:name="android.permission.DUMP" />
+    <uses-permission android:name="android.permission.PACKAGE_USAGE_STATS" />
+
+    <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
+                     android:targetPackage="com.android.internal.os.statsd.libstats"
+                     android:label="Tests for libstatspull">
+    </instrumentation>
+</manifest>
+
diff --git a/apex/tests/libstatspull/TEST_MAPPING b/apex/tests/libstatspull/TEST_MAPPING
new file mode 100644
index 0000000..5e1178c
--- /dev/null
+++ b/apex/tests/libstatspull/TEST_MAPPING
@@ -0,0 +1,7 @@
+{
+  "presubmit" : [
+    {
+      "name" : "LibStatsPullTests"
+    }
+  ]
+}
\ No newline at end of file
diff --git a/apex/tests/libstatspull/jni/stats_pull_helper.cpp b/apex/tests/libstatspull/jni/stats_pull_helper.cpp
new file mode 100644
index 0000000..e4ab823
--- /dev/null
+++ b/apex/tests/libstatspull/jni/stats_pull_helper.cpp
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2019, 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.
+ */
+
+#include <binder/ProcessState.h>
+#include <jni.h>
+#include <log/log.h>
+#include <stats_event.h>
+#include <stats_pull_atom_callback.h>
+
+#include <chrono>
+#include <thread>
+
+using std::this_thread::sleep_for;
+using namespace android;
+
+namespace {
+static int32_t sAtomTag;
+static int32_t sPullReturnVal;
+static int64_t sLatencyMillis;
+static int32_t sAtomsPerPull;
+static int32_t sNumPulls = 0;
+
+static bool initialized = false;
+
+static void init() {
+    if (!initialized) {
+        initialized = true;
+        // Set up the binder
+        sp<ProcessState> ps(ProcessState::self());
+        ps->setThreadPoolMaxThreadCount(9);
+        ps->startThreadPool();
+        ps->giveThreadPoolName();
+    }
+}
+
+static status_pull_atom_return_t pullAtomCallback(int32_t atomTag, pulled_stats_event_list* data,
+                                                  void* /*cookie*/) {
+    sNumPulls++;
+    sleep_for(std::chrono::milliseconds(sLatencyMillis));
+    for (int i = 0; i < sAtomsPerPull; i++) {
+        stats_event* event = add_stats_event_to_pull_data(data);
+        stats_event_set_atom_id(event, atomTag);
+        stats_event_write_int64(event, (int64_t) sNumPulls);
+        stats_event_build(event);
+    }
+    return sPullReturnVal;
+}
+
+extern "C"
+JNIEXPORT void JNICALL
+Java_com_android_internal_os_statsd_libstats_LibStatsPullTests_registerStatsPuller(
+        JNIEnv* /*env*/, jobject /* this */, jint atomTag, jlong timeoutNs, jlong coolDownNs,
+        jint pullRetVal, jlong latencyMillis, int atomsPerPull)
+{
+    init();
+    sAtomTag = atomTag;
+    sPullReturnVal = pullRetVal;
+    sLatencyMillis = latencyMillis;
+    sAtomsPerPull = atomsPerPull;
+    sNumPulls = 0;
+    pull_atom_metadata metadata = {.cool_down_ns = coolDownNs,
+                                   .timeout_ns = timeoutNs,
+                                   .additive_fields = nullptr,
+                                   .additive_fields_size = 0};
+    register_stats_pull_atom_callback(sAtomTag, &pullAtomCallback, &metadata, nullptr);
+}
+
+extern "C"
+JNIEXPORT void JNICALL
+Java_com_android_internal_os_statsd_libstats_LibStatsPullTests_unregisterStatsPuller(
+        JNIEnv* /*env*/, jobject /* this */, jint /*atomTag*/)
+{
+    unregister_stats_pull_atom_callback(sAtomTag);
+}
+} // namespace
\ No newline at end of file
diff --git a/apex/tests/libstatspull/protos/test_atoms.proto b/apex/tests/libstatspull/protos/test_atoms.proto
new file mode 100644
index 0000000..56c1b53
--- /dev/null
+++ b/apex/tests/libstatspull/protos/test_atoms.proto
@@ -0,0 +1,32 @@
+/*
+ * Copyright (C) 2020 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.
+ */
+syntax = "proto2";
+
+package com.android.internal.os.statsd.protos;
+
+option java_package = "com.android.internal.os.statsd.protos";
+option java_outer_classname = "TestAtoms";
+
+message PullCallbackAtomWrapper {
+  optional PullCallbackAtom pull_callback_atom = 150030;
+}
+
+message PullCallbackAtom {
+  optional int64 long_val = 1;
+}
+
+
+
diff --git a/apex/tests/libstatspull/src/com/android/internal/os/statsd/StatsConfigUtils.java b/apex/tests/libstatspull/src/com/android/internal/os/statsd/StatsConfigUtils.java
new file mode 100644
index 0000000..d0d1400
--- /dev/null
+++ b/apex/tests/libstatspull/src/com/android/internal/os/statsd/StatsConfigUtils.java
@@ -0,0 +1,124 @@
+/*
+ * Copyright (C) 2019 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.internal.os.statsd;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.app.StatsManager;
+import android.util.Log;
+
+import com.android.internal.os.StatsdConfigProto.AtomMatcher;
+import com.android.internal.os.StatsdConfigProto.FieldValueMatcher;
+import com.android.internal.os.StatsdConfigProto.SimpleAtomMatcher;
+import com.android.internal.os.StatsdConfigProto.StatsdConfig;
+import com.android.os.AtomsProto.AppBreadcrumbReported;
+import com.android.os.AtomsProto.Atom;
+import com.android.os.StatsLog.ConfigMetricsReport;
+import com.android.os.StatsLog.ConfigMetricsReportList;
+import com.android.os.StatsLog.GaugeBucketInfo;
+import com.android.os.StatsLog.GaugeMetricData;
+import com.android.os.StatsLog.StatsLogReport;
+import com.android.os.StatsLog.StatsdStatsReport;
+import com.android.os.StatsLog.StatsdStatsReport.ConfigStats;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Util class for constructing statsd configs.
+ */
+public class StatsConfigUtils {
+    public static final String TAG = "statsd.StatsConfigUtils";
+    public static final int SHORT_WAIT = 2_000; // 2 seconds.
+
+    /**
+     * @return An empty StatsdConfig in serialized proto format.
+     */
+    public static StatsdConfig.Builder getSimpleTestConfig(long configId) {
+        return StatsdConfig.newBuilder().setId(configId)
+                .addAllowedLogSource(StatsConfigUtils.class.getPackage().getName());
+    }
+
+
+    public static boolean verifyValidConfigExists(StatsManager statsManager, long configId) {
+        StatsdStatsReport report = null;
+        try {
+            report = StatsdStatsReport.parser().parseFrom(statsManager.getStatsMetadata());
+        } catch (Exception e) {
+            Log.e(TAG, "getMetadata failed", e);
+        }
+        assertThat(report).isNotNull();
+        boolean foundConfig = false;
+        for (ConfigStats configStats : report.getConfigStatsList()) {
+            if (configStats.getId() == configId && configStats.getIsValid()
+                    && configStats.getDeletionTimeSec() == 0) {
+                foundConfig = true;
+            }
+        }
+        return foundConfig;
+    }
+
+    public static AtomMatcher getAppBreadcrumbMatcher(long id, int label) {
+        return AtomMatcher.newBuilder()
+                .setId(id)
+                .setSimpleAtomMatcher(
+                        SimpleAtomMatcher.newBuilder()
+                                .setAtomId(Atom.APP_BREADCRUMB_REPORTED_FIELD_NUMBER)
+                                .addFieldValueMatcher(FieldValueMatcher.newBuilder()
+                                        .setField(AppBreadcrumbReported.LABEL_FIELD_NUMBER)
+                                        .setEqInt(label)
+                                )
+                )
+                .build();
+    }
+
+    public static ConfigMetricsReport getConfigMetricsReport(StatsManager statsManager,
+            long configId) {
+        ConfigMetricsReportList reportList = null;
+        try {
+            reportList = ConfigMetricsReportList.parser()
+                    .parseFrom(statsManager.getReports(configId));
+        } catch (Exception e) {
+            Log.e(TAG, "getData failed", e);
+        }
+        assertThat(reportList).isNotNull();
+        assertThat(reportList.getReportsCount()).isEqualTo(1);
+        ConfigMetricsReport report = reportList.getReports(0);
+        assertThat(report.getDumpReportReason())
+                .isEqualTo(ConfigMetricsReport.DumpReportReason.GET_DATA_CALLED);
+        return report;
+
+    }
+    public static List<Atom> getGaugeMetricDataList(ConfigMetricsReport report) {
+        List<Atom> data = new ArrayList<>();
+        for (StatsLogReport metric : report.getMetricsList()) {
+            for (GaugeMetricData gaugeMetricData : metric.getGaugeMetrics().getDataList()) {
+                for (GaugeBucketInfo bucketInfo : gaugeMetricData.getBucketInfoList()) {
+                    for (Atom atom : bucketInfo.getAtomList()) {
+                        data.add(atom);
+                    }
+                }
+            }
+        }
+        return data;
+    }
+
+    public static List<Atom> getGaugeMetricDataList(StatsManager statsManager, long configId) {
+        ConfigMetricsReport report = getConfigMetricsReport(statsManager, configId);
+        return getGaugeMetricDataList(report);
+    }
+}
+
diff --git a/apex/tests/libstatspull/src/com/android/internal/os/statsd/libstats/LibStatsPullTests.java b/apex/tests/libstatspull/src/com/android/internal/os/statsd/libstats/LibStatsPullTests.java
new file mode 100644
index 0000000..dbd636d
--- /dev/null
+++ b/apex/tests/libstatspull/src/com/android/internal/os/statsd/libstats/LibStatsPullTests.java
@@ -0,0 +1,285 @@
+/*
+ * Copyright (C) 2019 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.internal.os.statsd.libstats;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.app.StatsManager;
+import android.content.Context;
+import android.util.Log;
+import android.util.StatsLog;
+
+import androidx.test.InstrumentationRegistry;
+import androidx.test.filters.MediumTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.internal.os.StatsdConfigProto.AtomMatcher;
+import com.android.internal.os.StatsdConfigProto.FieldFilter;
+import com.android.internal.os.StatsdConfigProto.GaugeMetric;
+import com.android.internal.os.StatsdConfigProto.SimpleAtomMatcher;
+import com.android.internal.os.StatsdConfigProto.StatsdConfig;
+import com.android.internal.os.StatsdConfigProto.TimeUnit;
+import com.android.internal.os.statsd.StatsConfigUtils;
+import com.android.internal.os.statsd.protos.TestAtoms;
+import com.android.os.AtomsProto.Atom;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.List;
+
+/**
+ * Test puller registration.
+ */
+@MediumTest
+@RunWith(AndroidJUnit4.class)
+public class LibStatsPullTests {
+    private static final String LOG_TAG = LibStatsPullTests.class.getSimpleName();
+    private static final int SHORT_SLEEP_MILLIS = 250;
+    private static final int LONG_SLEEP_MILLIS = 1_000;
+    private Context mContext;
+    private static final int PULL_ATOM_TAG = 150030;
+    private static final int APP_BREADCRUMB_LABEL = 3;
+    private static int sPullReturnValue;
+    private static long sConfigId;
+    private static long sPullLatencyMillis;
+    private static long sPullTimeoutNs;
+    private static long sCoolDownNs;
+    private static int sAtomsPerPull;
+
+    static {
+        System.loadLibrary("statspull_testhelper");
+    }
+
+    /**
+     * Setup the tests. Initialize shared data.
+     */
+    @Before
+    public void setup() {
+//        Debug.waitForDebugger();
+        mContext = InstrumentationRegistry.getTargetContext();
+        assertThat(InstrumentationRegistry.getInstrumentation()).isNotNull();
+        sPullReturnValue = StatsManager.PULL_SUCCESS;
+        sPullLatencyMillis = 0;
+        sPullTimeoutNs = 10_000_000_000L;
+        sCoolDownNs = 1_000_000_000L;
+        sAtomsPerPull = 1;
+    }
+
+    /**
+     * Teardown the tests.
+     */
+    @After
+    public void tearDown() throws Exception {
+        unregisterStatsPuller(PULL_ATOM_TAG);
+        StatsManager statsManager = (StatsManager) mContext.getSystemService(
+                Context.STATS_MANAGER);
+        statsManager.removeConfig(sConfigId);
+    }
+
+    /**
+     * Tests adding a puller callback and that pulls complete successfully.
+     */
+    @Test
+    public void testPullAtomCallbackRegistration() throws Exception {
+        StatsManager statsManager = (StatsManager) mContext.getSystemService(
+                Context.STATS_MANAGER);
+        // Upload a config that captures that pulled atom.
+        createAndAddConfigToStatsd(statsManager);
+
+        // Add the puller.
+        registerStatsPuller(PULL_ATOM_TAG, sPullTimeoutNs, sCoolDownNs, sPullReturnValue,
+                sPullLatencyMillis, sAtomsPerPull);
+        Thread.sleep(SHORT_SLEEP_MILLIS);
+        StatsLog.logStart(APP_BREADCRUMB_LABEL);
+        // Let the current bucket finish.
+        Thread.sleep(LONG_SLEEP_MILLIS);
+        List<Atom> data = StatsConfigUtils.getGaugeMetricDataList(statsManager, sConfigId);
+        unregisterStatsPuller(PULL_ATOM_TAG);
+        assertThat(data.size()).isEqualTo(1);
+        TestAtoms.PullCallbackAtomWrapper atomWrapper = null;
+        try {
+            atomWrapper = TestAtoms.PullCallbackAtomWrapper.parser()
+                    .parseFrom(data.get(0).toByteArray());
+        } catch (Exception e) {
+            Log.e(LOG_TAG, "Failed to parse primitive atoms");
+        }
+        assertThat(atomWrapper).isNotNull();
+        assertThat(atomWrapper.hasPullCallbackAtom()).isTrue();
+        TestAtoms.PullCallbackAtom atom =
+                atomWrapper.getPullCallbackAtom();
+        assertThat(atom.getLongVal()).isEqualTo(1);
+    }
+
+    /**
+     * Tests that a failed pull is skipped.
+     */
+    @Test
+    public void testPullAtomCallbackFailure() throws Exception {
+        StatsManager statsManager = (StatsManager) mContext.getSystemService(
+                Context.STATS_MANAGER);
+        createAndAddConfigToStatsd(statsManager);
+        sPullReturnValue = StatsManager.PULL_SKIP;
+        // Add the puller.
+        registerStatsPuller(PULL_ATOM_TAG, sPullTimeoutNs, sCoolDownNs, sPullReturnValue,
+                sPullLatencyMillis, sAtomsPerPull);
+        Thread.sleep(SHORT_SLEEP_MILLIS);
+        StatsLog.logStart(APP_BREADCRUMB_LABEL);
+        // Let the current bucket finish.
+        Thread.sleep(LONG_SLEEP_MILLIS);
+        List<Atom> data = StatsConfigUtils.getGaugeMetricDataList(statsManager, sConfigId);
+        unregisterStatsPuller(PULL_ATOM_TAG);
+        assertThat(data.size()).isEqualTo(0);
+    }
+
+    /**
+     * Tests that a pull that times out is skipped.
+     */
+    @Test
+    public void testPullAtomCallbackTimeout() throws Exception {
+        StatsManager statsManager = (StatsManager) mContext.getSystemService(
+                Context.STATS_MANAGER);
+        createAndAddConfigToStatsd(statsManager);
+        // The puller will sleep for 1.5 sec.
+        sPullLatencyMillis = 1_500;
+        // 1 second timeout
+        sPullTimeoutNs = 1_000_000_000;
+
+        // Add the puller.
+        registerStatsPuller(PULL_ATOM_TAG, sPullTimeoutNs, sCoolDownNs, sPullReturnValue,
+                sPullLatencyMillis, sAtomsPerPull);
+        Thread.sleep(SHORT_SLEEP_MILLIS);
+        StatsLog.logStart(APP_BREADCRUMB_LABEL);
+        // Let the current bucket finish and the pull timeout.
+        Thread.sleep(sPullLatencyMillis * 2);
+        List<Atom> data = StatsConfigUtils.getGaugeMetricDataList(statsManager, sConfigId);
+        unregisterStatsPuller(PULL_ATOM_TAG);
+        assertThat(data.size()).isEqualTo(0);
+    }
+
+    /**
+     * Tests that 2 pulls in quick succession use the cache instead of pulling again.
+     */
+    @Test
+    public void testPullAtomCallbackCache() throws Exception {
+        StatsManager statsManager = (StatsManager) mContext.getSystemService(
+                Context.STATS_MANAGER);
+        createAndAddConfigToStatsd(statsManager);
+
+        // Set the cooldown to 10 seconds
+        sCoolDownNs = 10_000_000_000L;
+        // Add the puller.
+        registerStatsPuller(PULL_ATOM_TAG, sPullTimeoutNs, sCoolDownNs, sPullReturnValue,
+                sPullLatencyMillis, sAtomsPerPull);
+
+        Thread.sleep(SHORT_SLEEP_MILLIS);
+        StatsLog.logStart(APP_BREADCRUMB_LABEL);
+        // Pull from cache.
+        StatsLog.logStart(APP_BREADCRUMB_LABEL);
+        Thread.sleep(LONG_SLEEP_MILLIS);
+        List<Atom> data = StatsConfigUtils.getGaugeMetricDataList(statsManager, sConfigId);
+        unregisterStatsPuller(PULL_ATOM_TAG);
+        assertThat(data.size()).isEqualTo(2);
+        for (int i = 0; i < data.size(); i++) {
+            TestAtoms.PullCallbackAtomWrapper atomWrapper = null;
+            try {
+                atomWrapper = TestAtoms.PullCallbackAtomWrapper.parser()
+                        .parseFrom(data.get(i).toByteArray());
+            } catch (Exception e) {
+                Log.e(LOG_TAG, "Failed to parse primitive atoms");
+            }
+            assertThat(atomWrapper).isNotNull();
+            assertThat(atomWrapper.hasPullCallbackAtom()).isTrue();
+            TestAtoms.PullCallbackAtom atom =
+                    atomWrapper.getPullCallbackAtom();
+            assertThat(atom.getLongVal()).isEqualTo(1);
+        }
+    }
+
+    /**
+     * Tests that a pull that returns 1000 stats events works properly.
+     */
+    @Test
+    public void testPullAtomCallbackStress() throws Exception {
+        StatsManager statsManager = (StatsManager) mContext.getSystemService(
+                Context.STATS_MANAGER);
+        // Upload a config that captures that pulled atom.
+        createAndAddConfigToStatsd(statsManager);
+        sAtomsPerPull = 1000;
+        // Add the puller.
+        registerStatsPuller(PULL_ATOM_TAG, sPullTimeoutNs, sCoolDownNs, sPullReturnValue,
+                sPullLatencyMillis, sAtomsPerPull);
+
+        Thread.sleep(SHORT_SLEEP_MILLIS);
+        StatsLog.logStart(APP_BREADCRUMB_LABEL);
+        // Let the current bucket finish.
+        Thread.sleep(LONG_SLEEP_MILLIS);
+        List<Atom> data = StatsConfigUtils.getGaugeMetricDataList(statsManager, sConfigId);
+        statsManager.unregisterPullAtomCallback(PULL_ATOM_TAG);
+        assertThat(data.size()).isEqualTo(sAtomsPerPull);
+
+        for (int i = 0; i < data.size(); i++) {
+            TestAtoms.PullCallbackAtomWrapper atomWrapper = null;
+            try {
+                atomWrapper = TestAtoms.PullCallbackAtomWrapper.parser()
+                        .parseFrom(data.get(i).toByteArray());
+            } catch (Exception e) {
+                Log.e(LOG_TAG, "Failed to parse primitive atoms");
+            }
+            assertThat(atomWrapper).isNotNull();
+            assertThat(atomWrapper.hasPullCallbackAtom()).isTrue();
+            TestAtoms.PullCallbackAtom atom =
+                    atomWrapper.getPullCallbackAtom();
+            assertThat(atom.getLongVal()).isEqualTo(1);
+        }
+    }
+
+    private void createAndAddConfigToStatsd(StatsManager statsManager) throws Exception {
+        sConfigId = System.currentTimeMillis();
+        long triggerMatcherId = sConfigId + 10;
+        long pullerMatcherId = sConfigId + 11;
+        long metricId = sConfigId + 100;
+        StatsdConfig config = StatsConfigUtils.getSimpleTestConfig(sConfigId)
+                .addAtomMatcher(
+                        StatsConfigUtils.getAppBreadcrumbMatcher(triggerMatcherId,
+                                APP_BREADCRUMB_LABEL))
+                .addAtomMatcher(AtomMatcher.newBuilder()
+                        .setId(pullerMatcherId)
+                        .setSimpleAtomMatcher(SimpleAtomMatcher.newBuilder()
+                                .setAtomId(PULL_ATOM_TAG))
+                )
+                .addGaugeMetric(GaugeMetric.newBuilder()
+                        .setId(metricId)
+                        .setWhat(pullerMatcherId)
+                        .setTriggerEvent(triggerMatcherId)
+                        .setGaugeFieldsFilter(FieldFilter.newBuilder().setIncludeAll(true))
+                        .setBucket(TimeUnit.CTS)
+                        .setSamplingType(GaugeMetric.SamplingType.FIRST_N_SAMPLES)
+                        .setMaxNumGaugeAtomsPerBucket(1000)
+                )
+                .build();
+        statsManager.addConfig(sConfigId, config.toByteArray());
+        assertThat(StatsConfigUtils.verifyValidConfigExists(statsManager, sConfigId)).isTrue();
+    }
+
+    private native void registerStatsPuller(int atomTag, long timeoutNs, long coolDownNs,
+            int pullReturnVal, long latencyMillis, int atomPerPull);
+
+    private native void unregisterStatsPuller(int atomTag);
+}
+