path: root/statsd/tests/shell/ShellSubscriber_test.cpp

// Copyright (C) 2018 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 "src/shell/ShellSubscriber.h"

#include <aidl/android/os/StatsSubscriptionCallbackReason.h>
#include <gtest/gtest.h>
#include <stdio.h>
#include <unistd.h>

#include <optional>
#include <vector>

#include "frameworks/proto_logging/stats/atoms.pb.h"
#include "gtest_matchers.h"
#include "src/shell/shell_config.pb.h"
#include "src/shell/shell_data.pb.h"
#include "stats_event.h"
#include "statslog_statsdtest.h"
#include "tests/metrics/metrics_test_helper.h"
#include "tests/statsd_test_util.h"

using ::aidl::android::os::StatsSubscriptionCallbackReason;
using android::sp;
using android::os::statsd::TestAtomReported;
using android::os::statsd::TrainExperimentIds;
using android::os::statsd::util::BytesField;
using android::os::statsd::util::CPU_ACTIVE_TIME;
using android::os::statsd::util::PHONE_SIGNAL_STRENGTH_CHANGED;
using android::os::statsd::util::PLUGGED_STATE_CHANGED;
using android::os::statsd::util::SCREEN_STATE_CHANGED;
using android::os::statsd::util::TEST_ATOM_REPORTED;
using std::vector;
using testing::_;
using testing::A;
using testing::ByMove;
using testing::DoAll;
using testing::Invoke;
using testing::NaggyMock;
using testing::Return;
using testing::SaveArg;
using testing::SetArgPointee;
using testing::StrictMock;

namespace android {
namespace os {
namespace statsd {

#ifdef __ANDROID__

namespace {

int kUid1 = 1000;
int kUid2 = 2000;

int kCpuTime1 = 100;
int kCpuTime2 = 200;

int64_t kCpuActiveTimeEventTimestampNs = 1111L;

// Number of clients running simultaneously

// Just a single client
const int kSingleClient = 1;
// One more client than allowed binder threads
const int kNumClients = 11;

// Utility to make an expected pulled atom shell data
ShellData getExpectedPulledData() {
    ShellData shellData;
    auto* atom1 = shellData.add_atom()->mutable_cpu_active_time();
    atom1->set_uid(kUid1);
    atom1->set_time_millis(kCpuTime1);
    shellData.add_elapsed_timestamp_nanos(kCpuActiveTimeEventTimestampNs);

    auto* atom2 = shellData.add_atom()->mutable_cpu_active_time();
    atom2->set_uid(kUid2);
    atom2->set_time_millis(kCpuTime2);
    shellData.add_elapsed_timestamp_nanos(kCpuActiveTimeEventTimestampNs);

    return shellData;
}

// Utility to make a pulled atom Shell Config
ShellSubscription getPulledConfig() {
    ShellSubscription config;
    auto* pull_config = config.add_pulled();
    pull_config->mutable_matcher()->set_atom_id(CPU_ACTIVE_TIME);
    pull_config->set_freq_millis(2000);
    return config;
}

// Utility to adjust CPU time for pulled events
shared_ptr<LogEvent> makeCpuActiveTimeAtom(int32_t uid, int64_t timeMillis) {
    AStatsEvent* statsEvent = AStatsEvent_obtain();
    AStatsEvent_setAtomId(statsEvent, CPU_ACTIVE_TIME);
    AStatsEvent_overwriteTimestamp(statsEvent, kCpuActiveTimeEventTimestampNs);
    AStatsEvent_writeInt32(statsEvent, uid);
    AStatsEvent_writeInt64(statsEvent, timeMillis);

    std::shared_ptr<LogEvent> logEvent = std::make_shared<LogEvent>(/*uid=*/0, /*pid=*/0);
    parseStatsEventToLogEvent(statsEvent, logEvent.get());
    return logEvent;
}

// Utility to create pushed atom LogEvents
vector<std::shared_ptr<LogEvent>> getPushedEvents() {
    vector<std::shared_ptr<LogEvent>> pushedList;
    // Create the LogEvent from an AStatsEvent
    std::unique_ptr<LogEvent> logEvent1 = CreateScreenStateChangedEvent(
            1000 /*timestamp*/, ::android::view::DisplayStateEnum::DISPLAY_STATE_ON);
    std::unique_ptr<LogEvent> logEvent2 = CreateScreenStateChangedEvent(
            2000 /*timestamp*/, ::android::view::DisplayStateEnum::DISPLAY_STATE_OFF);
    std::unique_ptr<LogEvent> logEvent3 = CreateBatteryStateChangedEvent(
            3000 /*timestamp*/, BatteryPluggedStateEnum::BATTERY_PLUGGED_USB);
    std::unique_ptr<LogEvent> logEvent4 = CreateBatteryStateChangedEvent(
            4000 /*timestamp*/, BatteryPluggedStateEnum::BATTERY_PLUGGED_NONE);
    pushedList.push_back(std::move(logEvent1));
    pushedList.push_back(std::move(logEvent2));
    pushedList.push_back(std::move(logEvent3));
    pushedList.push_back(std::move(logEvent4));
    return pushedList;
}

// Utility to read & return ShellData proto, skipping heartbeats.
static ShellData readData(int fd) {
    ssize_t dataSize = 0;
    while (dataSize == 0) {
        read(fd, &dataSize, sizeof(dataSize));
    }
    // Read that much data in proto binary format.
    vector<uint8_t> dataBuffer(dataSize);
    EXPECT_EQ((int)dataSize, read(fd, dataBuffer.data(), dataSize));

    // Make sure the received bytes can be parsed to an atom.
    ShellData receivedAtom;
    EXPECT_TRUE(receivedAtom.ParseFromArray(dataBuffer.data(), dataSize) != 0);
    return receivedAtom;
}

void runShellTest(ShellSubscription config, sp<MockUidMap> uidMap,
                  sp<MockStatsPullerManager> pullerManager,
                  const vector<std::shared_ptr<LogEvent>>& pushedEvents,
                  const vector<ShellData>& expectedData, int numClients) {
    sp<ShellSubscriber> shellManager = new ShellSubscriber(uidMap, pullerManager);

    size_t bufferSize = config.ByteSize();
    vector<uint8_t> buffer(bufferSize);
    config.SerializeToArray(&buffer[0], bufferSize);

    int fds_configs[numClients][2];
    int fds_datas[numClients][2];
    for (int i = 0; i < numClients; i++) {
        // set up 2 pipes for read/write config and data
        ASSERT_EQ(0, pipe2(fds_configs[i], O_CLOEXEC));
        ASSERT_EQ(0, pipe2(fds_datas[i], O_CLOEXEC));

        // write the config to pipe, first write size of the config
        write(fds_configs[i][1], &bufferSize, sizeof(bufferSize));
        // then write config itself
        write(fds_configs[i][1], buffer.data(), bufferSize);
        close(fds_configs[i][1]);

        shellManager->startNewSubscription(fds_configs[i][0], fds_datas[i][1], /*timeoutSec=*/-1);
        close(fds_configs[i][0]);
        close(fds_datas[i][1]);
    }

    // send a log event that matches the config.
    for (const auto& event : pushedEvents) {
        shellManager->onLogEvent(*event);
    }

    for (int i = 0; i < numClients; i++) {
        vector<ShellData> actualData;
        for (int j = 1; j <= expectedData.size(); j++) {
            actualData.push_back(readData(fds_datas[i][0]));
        }

        EXPECT_THAT(expectedData, UnorderedPointwise(EqShellData(), actualData));
    }

    // Not closing fds_datas[i][0] because this causes writes within ShellSubscriberClient to hang
}

unique_ptr<LogEvent> createTestAtomReportedEvent(const uint64_t timestampNs,
                                                 const int32_t intFieldValue,
                                                 const vector<int64_t>& expIds) {
    TrainExperimentIds trainExpIds;
    *trainExpIds.mutable_experiment_id() = {expIds.begin(), expIds.end()};
    const vector<uint8_t> trainExpIdsBytes = protoToBytes(trainExpIds);
    return CreateTestAtomReportedEvent(
            timestampNs, /* attributionUids */ {1001},
            /* attributionTags */ {"app1"}, intFieldValue, /*longField */ 0LL,
            /* floatField */ 0.0f,
            /* stringField */ "abc", /* boolField */ false, TestAtomReported::OFF, trainExpIdsBytes,
            /* repeatedIntField */ {}, /* repeatedLongField */ {}, /* repeatedFloatField */ {},
            /* repeatedStringField */ {}, /* repeatedBoolField */ {},
            /* repeatedBoolFieldLength */ 0, /* repeatedEnumField */ {});
}

TestAtomReported createTestAtomReportedProto(const int32_t intFieldValue,
                                             const vector<int64_t>& expIds) {
    TestAtomReported t;
    auto* attributionNode = t.add_attribution_node();
    attributionNode->set_uid(1001);
    attributionNode->set_tag("app1");
    t.set_int_field(intFieldValue);
    t.set_long_field(0);
    t.set_float_field(0.0f);
    t.set_string_field("abc");
    t.set_boolean_field(false);
    t.set_state(TestAtomReported_State_OFF);
    *t.mutable_bytes_field()->mutable_experiment_id() = {expIds.begin(), expIds.end()};
    return t;
}

class ShellSubscriberCallbackTest : public ::testing::Test {
protected:
    ShellSubscriberCallbackTest()
        : uidMap(new NaggyMock<MockUidMap>()),
          pullerManager(new StrictMock<MockStatsPullerManager>()),
          shellSubscriber(uidMap, pullerManager),
          callback(SharedRefBase::make<StrictMock<MockStatsSubscriptionCallback>>()),
          reason(nullopt) {
    }

    void SetUp() override {
        // Save callback arguments when it is invoked.
        ON_CALL(*callback, onSubscriptionData(_, _))
                .WillByDefault(DoAll(SaveArg<0>(&reason), SaveArg<1>(&payload),
                                     Return(ByMove(Status::ok()))));

        ShellSubscription config;
        config.add_pushed()->set_atom_id(TEST_ATOM_REPORTED);
        config.add_pushed()->set_atom_id(SCREEN_STATE_CHANGED);
        config.add_pushed()->set_atom_id(PHONE_SIGNAL_STRENGTH_CHANGED);
        configBytes = protoToBytes(config);
    }

    sp<MockUidMap> uidMap;
    sp<MockStatsPullerManager> pullerManager;
    ShellSubscriber shellSubscriber;
    std::shared_ptr<MockStatsSubscriptionCallback> callback;
    vector<uint8_t> configBytes;

    // Capture callback arguments.
    std::optional<StatsSubscriptionCallbackReason> reason;
    vector<uint8_t> payload;
};

class ShellSubscriberCallbackPulledTest : public ShellSubscriberCallbackTest {
protected:
    void SetUp() override {
        ShellSubscriberCallbackTest::SetUp();

        const vector<int32_t> uids{AID_SYSTEM};
        const vector<std::shared_ptr<LogEvent>> pulledData{
                makeCpuActiveTimeAtom(/*uid=*/kUid1, /*timeMillis=*/kCpuTime1),
                makeCpuActiveTimeAtom(/*uid=*/kUid2, /*timeMillis=*/kCpuTime2)};
        ON_CALL(*pullerManager, Pull(CPU_ACTIVE_TIME, uids, _, _))
                .WillByDefault(DoAll(SetArgPointee<3>(pulledData), Return(true)));

        configBytes = protoToBytes(getPulledConfig());

        // Used to call pullAndSendHeartbeatsIfNeeded directly without depending on sleep.
        shellSubscriberClient = std::move(ShellSubscriberClient::create(
                configBytes, callback, /* startTimeSec= */ 0, uidMap, pullerManager));
    }

    unique_ptr<ShellSubscriberClient> shellSubscriberClient;
};

}  // namespace

TEST_F(ShellSubscriberCallbackTest, testAddSubscription) {
    EXPECT_TRUE(shellSubscriber.startNewSubscription(configBytes, callback));
}

TEST_F(ShellSubscriberCallbackTest, testAddSubscriptionExceedMax) {
    const size_t maxSubs = ShellSubscriber::getMaxSubscriptions();
    vector<bool> results(maxSubs, false);
    for (int i = 0; i < maxSubs; i++) {
        results[i] = shellSubscriber.startNewSubscription(configBytes, callback);
    }

    // First maxSubs subscriptions should succeed.
    EXPECT_THAT(results, Each(IsTrue()));

    // Subsequent startNewSubscription should fail.
    EXPECT_FALSE(shellSubscriber.startNewSubscription(configBytes, callback));
}

TEST_F(ShellSubscriberCallbackTest, testPushedEventsAreCached) {
    // Expect callback to not be invoked
    EXPECT_CALL(*callback, onSubscriptionData(_, _)).Times(Exactly(0));

    shellSubscriber.startNewSubscription(configBytes, callback);

    // Log an event that does NOT invoke the callack.
    shellSubscriber.onLogEvent(*CreateScreenStateChangedEvent(
            1000 /*timestamp*/, ::android::view::DisplayStateEnum::DISPLAY_STATE_ON));
}

TEST_F(ShellSubscriberCallbackTest, testOverflowCacheIsFlushed) {
    // Expect callback to be invoked once.
    EXPECT_CALL(*callback, onSubscriptionData(_, _)).Times(Exactly(1));

    shellSubscriber.startNewSubscription(configBytes, callback);

    shellSubscriber.onLogEvent(*CreateScreenStateChangedEvent(
            1000 /*timestamp*/, ::android::view::DisplayStateEnum::DISPLAY_STATE_ON));

    // Inflate size of TestAtomReported through the MODE_BYTES field.
    const vector<int64_t> expIds = vector<int64_t>(200, INT64_MAX);

    // This event should trigger cache overflow flush.
    shellSubscriber.onLogEvent(*createTestAtomReportedEvent(/*timestampNs=*/1100,
                                                            /*intFieldValue=*/1, expIds));

    EXPECT_THAT(reason, Eq(StatsSubscriptionCallbackReason::STATSD_INITIATED));

    // Get ShellData proto from the bytes payload of the callback.
    ShellData actualShellData;
    ASSERT_TRUE(actualShellData.ParseFromArray(payload.data(), payload.size()));

    ShellData expectedShellData;
    expectedShellData.add_atom()->mutable_screen_state_changed()->set_state(
            ::android::view::DisplayStateEnum::DISPLAY_STATE_ON);
    *expectedShellData.add_atom()->mutable_test_atom_reported() =
            createTestAtomReportedProto(/* intFieldValue=*/1, expIds);
    expectedShellData.add_elapsed_timestamp_nanos(1000);
    expectedShellData.add_elapsed_timestamp_nanos(1100);

    EXPECT_THAT(actualShellData, EqShellData(expectedShellData));
}

TEST_F(ShellSubscriberCallbackTest, testFlushTrigger) {
    // Expect callback to be invoked once.
    EXPECT_CALL(*callback, onSubscriptionData(_, _)).Times(Exactly(1));

    shellSubscriber.startNewSubscription(configBytes, callback);

    shellSubscriber.onLogEvent(*CreateScreenStateChangedEvent(
            1000 /*timestamp*/, ::android::view::DisplayStateEnum::DISPLAY_STATE_ON));

    shellSubscriber.flushSubscription(callback);

    EXPECT_THAT(reason, Eq(StatsSubscriptionCallbackReason::FLUSH_REQUESTED));

    // Get ShellData proto from the bytes payload of the callback.
    ShellData actualShellData;
    ASSERT_TRUE(actualShellData.ParseFromArray(payload.data(), payload.size()));

    ShellData expectedShellData;
    expectedShellData.add_atom()->mutable_screen_state_changed()->set_state(
            ::android::view::DisplayStateEnum::DISPLAY_STATE_ON);
    expectedShellData.add_elapsed_timestamp_nanos(1000);

    EXPECT_THAT(actualShellData, EqShellData(expectedShellData));
}

TEST_F(ShellSubscriberCallbackTest, testFlushTriggerEmptyCache) {
    // Expect callback to be invoked once.
    EXPECT_CALL(*callback, onSubscriptionData(_, _)).Times(Exactly(1));

    shellSubscriber.startNewSubscription(configBytes, callback);

    shellSubscriber.flushSubscription(callback);

    EXPECT_THAT(reason, Eq(StatsSubscriptionCallbackReason::FLUSH_REQUESTED));

    // Get ShellData proto from the bytes payload of the callback.
    ShellData actualShellData;
    ASSERT_TRUE(actualShellData.ParseFromArray(payload.data(), payload.size()));

    ShellData expectedShellData;

    EXPECT_THAT(actualShellData, EqShellData(expectedShellData));
}

TEST_F(ShellSubscriberCallbackTest, testUnsubscribe) {
    // Expect callback to be invoked once.
    EXPECT_CALL(*callback, onSubscriptionData(_, _)).Times(Exactly(1));

    shellSubscriber.startNewSubscription(configBytes, callback);

    shellSubscriber.onLogEvent(*CreateScreenStateChangedEvent(
            1000 /*timestamp*/, ::android::view::DisplayStateEnum::DISPLAY_STATE_ON));

    shellSubscriber.unsubscribe(callback);

    EXPECT_THAT(reason, Eq(StatsSubscriptionCallbackReason::SUBSCRIPTION_ENDED));

    // Get ShellData proto from the bytes payload of the callback.
    ShellData actualShellData;
    ASSERT_TRUE(actualShellData.ParseFromArray(payload.data(), payload.size()));

    ShellData expectedShellData;
    expectedShellData.add_atom()->mutable_screen_state_changed()->set_state(
            ::android::view::DisplayStateEnum::DISPLAY_STATE_ON);
    expectedShellData.add_elapsed_timestamp_nanos(1000);

    EXPECT_THAT(actualShellData, EqShellData(expectedShellData));

    // This event is ignored as the subscription has ended.
    shellSubscriber.onLogEvent(*CreateScreenStateChangedEvent(
            1000 /*timestamp*/, ::android::view::DisplayStateEnum::DISPLAY_STATE_ON));

    // This should be a no-op as we've already unsubscribed.
    shellSubscriber.unsubscribe(callback);
}

TEST_F(ShellSubscriberCallbackTest, testUnsubscribeEmptyCache) {
    // Expect callback to be invoked once.
    EXPECT_CALL(*callback, onSubscriptionData(_, _)).Times(Exactly(1));

    shellSubscriber.startNewSubscription(configBytes, callback);

    shellSubscriber.unsubscribe(callback);

    EXPECT_THAT(reason, Eq(StatsSubscriptionCallbackReason::SUBSCRIPTION_ENDED));

    // Get ShellData proto from the bytes payload of the callback.
    ShellData actualShellData;
    ASSERT_TRUE(actualShellData.ParseFromArray(payload.data(), payload.size()));

    ShellData expectedShellData;

    EXPECT_THAT(actualShellData, EqShellData(expectedShellData));
}

TEST_F(ShellSubscriberCallbackTest, testTruncateTimestampAtom) {
    // Expect callback to be invoked once.
    EXPECT_CALL(*callback, onSubscriptionData(_, _)).Times(Exactly(1));

    shellSubscriber.startNewSubscription(configBytes, callback);

    shellSubscriber.onLogEvent(*CreatePhoneSignalStrengthChangedEvent(
            NS_PER_SEC * 5 * 60 + 1000 /*timestamp*/,
            ::android::telephony::SignalStrengthEnum::SIGNAL_STRENGTH_GOOD));

    shellSubscriber.flushSubscription(callback);

    // Get ShellData proto from the bytes payload of the callback.
    ShellData actualShellData;
    ASSERT_TRUE(actualShellData.ParseFromArray(payload.data(), payload.size()));

    ShellData expectedShellData;
    expectedShellData.add_atom()->mutable_phone_signal_strength_changed()->set_signal_strength(
            ::android::telephony::SignalStrengthEnum::SIGNAL_STRENGTH_GOOD);
    expectedShellData.add_elapsed_timestamp_nanos(NS_PER_SEC * 5 * 60);

    EXPECT_THAT(actualShellData, EqShellData(expectedShellData));
}

TEST_F(ShellSubscriberCallbackPulledTest, testPullIfNeededBeforeInterval) {
    // Pull should not happen
    EXPECT_CALL(*pullerManager, Pull(_, A<const vector<int32_t>&>(), _, _)).Times(Exactly(0));

    // Expect callback to not be invoked.
    EXPECT_CALL(*callback, onSubscriptionData(_, _)).Times(Exactly(0));

    shellSubscriberClient->pullAndSendHeartbeatsIfNeeded(/* nowSecs= */ 0, /* nowMillis= */ 0,
                                                         /* nowNanos= */ 0);
}

TEST_F(ShellSubscriberCallbackPulledTest, testPullAtInterval) {
    // Pull should happen once. The data is cached.
    EXPECT_CALL(*pullerManager, Pull(_, A<const vector<int32_t>&>(), _, _)).Times(Exactly(1));

    // Expect callback to not be invoked.
    EXPECT_CALL(*callback, onSubscriptionData(_, _)).Times(Exactly(0));

    // This pull should NOT trigger a cache flush.
    shellSubscriberClient->pullAndSendHeartbeatsIfNeeded(/* nowSecs= */ 3, /* nowMillis= */ 3000,
                                                         /* nowNanos= */ 3'000'000'000);
}

TEST_F(ShellSubscriberCallbackPulledTest, testCachedPullIsFlushed) {
    // Pull should happen once. The data is cached.
    EXPECT_CALL(*pullerManager, Pull(_, A<const vector<int32_t>&>(), _, _)).Times(Exactly(1));

    // This pull should NOT trigger a cache flush.
    shellSubscriberClient->pullAndSendHeartbeatsIfNeeded(/* nowSecs= */ 3, /* nowMillis= */ 3000,
                                                         /* nowNanos= */ 3'000'000'000);

    // Expect callback to be invoked once flush is requested.
    EXPECT_CALL(*callback, onSubscriptionData(_, _)).Times(Exactly(1));

    // This should flush out data cached from the pull.
    shellSubscriberClient->flush();

    EXPECT_THAT(reason, Eq(StatsSubscriptionCallbackReason::FLUSH_REQUESTED));

    // Get ShellData proto from the bytes payload of the callback.
    ShellData actualShellData;
    ASSERT_TRUE(actualShellData.ParseFromArray(payload.data(), payload.size()));

    EXPECT_THAT(actualShellData, EqShellData(getExpectedPulledData()));
}

TEST_F(ShellSubscriberCallbackPulledTest, testPullAtCacheTimeout) {
    // Pull should happen once. The data is flushed.
    EXPECT_CALL(*pullerManager, Pull(_, A<const vector<int32_t>&>(), _, _)).Times(Exactly(1));

    // Expect callback to be invoked.
    EXPECT_CALL(*callback, onSubscriptionData(_, _)).Times(Exactly(1));

    // This pull should trigger a cache flush.
    shellSubscriberClient->pullAndSendHeartbeatsIfNeeded(/* nowSecs= */ 4, /* nowMillis= */ 4000,
                                                         /* nowNanos= */ 4'000'000'000);

    EXPECT_THAT(reason, Eq(StatsSubscriptionCallbackReason::STATSD_INITIATED));

    // Get ShellData proto from the bytes payload of the callback.
    ShellData actualShellData;
    ASSERT_TRUE(actualShellData.ParseFromArray(payload.data(), payload.size()));

    EXPECT_THAT(actualShellData, EqShellData(getExpectedPulledData()));
}

TEST(ShellSubscriberTest, testPushedSubscription) {
    sp<MockUidMap> uidMap = new NaggyMock<MockUidMap>();
    sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();

    vector<std::shared_ptr<LogEvent>> pushedList = getPushedEvents();

    // create a simple config to get screen events
    ShellSubscription config;
    config.add_pushed()->set_atom_id(SCREEN_STATE_CHANGED);

    // this is the expected screen event atom.
    vector<ShellData> expectedData;
    ShellData shellData1;
    shellData1.add_atom()->mutable_screen_state_changed()->set_state(
            ::android::view::DisplayStateEnum::DISPLAY_STATE_ON);
    shellData1.add_elapsed_timestamp_nanos(pushedList[0]->GetElapsedTimestampNs());
    ShellData shellData2;
    shellData2.add_atom()->mutable_screen_state_changed()->set_state(
            ::android::view::DisplayStateEnum::DISPLAY_STATE_OFF);
    shellData2.add_elapsed_timestamp_nanos(pushedList[1]->GetElapsedTimestampNs());
    expectedData.push_back(shellData1);
    expectedData.push_back(shellData2);

    // Test with single client
    TRACE_CALL(runShellTest, config, uidMap, pullerManager, pushedList, expectedData,
               kSingleClient);

    // Test with multiple client
    TRACE_CALL(runShellTest, config, uidMap, pullerManager, pushedList, expectedData, kNumClients);
}

TEST(ShellSubscriberTest, testPulledSubscription) {
    sp<MockUidMap> uidMap = new NaggyMock<MockUidMap>();
    sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();

    const vector<int32_t> uids = {AID_SYSTEM};
    EXPECT_CALL(*pullerManager, Pull(CPU_ACTIVE_TIME, uids, _, _))
            .WillRepeatedly(Invoke([](int tagId, const vector<int32_t>&, const int64_t,
                                      vector<std::shared_ptr<LogEvent>>* data) {
                data->clear();
                data->push_back(makeCpuActiveTimeAtom(/*uid=*/kUid1, /*timeMillis=*/kCpuTime1));
                data->push_back(makeCpuActiveTimeAtom(/*uid=*/kUid2, /*timeMillis=*/kCpuTime2));
                return true;
            }));

    // Test with single client
    TRACE_CALL(runShellTest, getPulledConfig(), uidMap, pullerManager, /*pushedEvents=*/{},
               {getExpectedPulledData()}, kSingleClient);

    // Test with multiple clients.
    TRACE_CALL(runShellTest, getPulledConfig(), uidMap, pullerManager, {},
               {getExpectedPulledData()}, kNumClients);
}

TEST(ShellSubscriberTest, testBothSubscriptions) {
    sp<MockUidMap> uidMap = new NaggyMock<MockUidMap>();
    sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();

    const vector<int32_t> uids = {AID_SYSTEM};
    EXPECT_CALL(*pullerManager, Pull(CPU_ACTIVE_TIME, uids, _, _))
            .WillRepeatedly(Invoke([](int tagId, const vector<int32_t>&, const int64_t,
                                      vector<std::shared_ptr<LogEvent>>* data) {
                data->clear();
                data->push_back(makeCpuActiveTimeAtom(/*uid=*/kUid1, /*timeMillis=*/kCpuTime1));
                data->push_back(makeCpuActiveTimeAtom(/*uid=*/kUid2, /*timeMillis=*/kCpuTime2));
                return true;
            }));

    vector<std::shared_ptr<LogEvent>> pushedList = getPushedEvents();

    ShellSubscription config = getPulledConfig();
    config.add_pushed()->set_atom_id(SCREEN_STATE_CHANGED);

    vector<ShellData> expectedData;
    ShellData shellData1;
    shellData1.add_atom()->mutable_screen_state_changed()->set_state(
            ::android::view::DisplayStateEnum::DISPLAY_STATE_ON);
    shellData1.add_elapsed_timestamp_nanos(pushedList[0]->GetElapsedTimestampNs());
    ShellData shellData2;
    shellData2.add_atom()->mutable_screen_state_changed()->set_state(
            ::android::view::DisplayStateEnum::DISPLAY_STATE_OFF);
    shellData2.add_elapsed_timestamp_nanos(pushedList[1]->GetElapsedTimestampNs());
    expectedData.push_back(getExpectedPulledData());
    expectedData.push_back(shellData1);
    expectedData.push_back(shellData2);

    // Test with single client
    TRACE_CALL(runShellTest, config, uidMap, pullerManager, pushedList, expectedData,
               kSingleClient);

    // Test with multiple client
    TRACE_CALL(runShellTest, config, uidMap, pullerManager, pushedList, expectedData, kNumClients);
}

TEST(ShellSubscriberTest, testMaxSizeGuard) {
    sp<MockUidMap> uidMap = new NaggyMock<MockUidMap>();
    sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
    sp<ShellSubscriber> shellManager = new ShellSubscriber(uidMap, pullerManager);

    // set up 2 pipes for read/write config and data
    int fds_config[2];
    ASSERT_EQ(0, pipe2(fds_config, O_CLOEXEC));

    int fds_data[2];
    ASSERT_EQ(0, pipe2(fds_data, O_CLOEXEC));

    // write invalid size of the config
    size_t invalidBufferSize = (shellManager->getMaxSizeKb() * 1024) + 1;
    write(fds_config[1], &invalidBufferSize, sizeof(invalidBufferSize));
    close(fds_config[1]);
    close(fds_data[0]);

    EXPECT_FALSE(shellManager->startNewSubscription(fds_config[0], fds_data[1], /*timeoutSec=*/-1));
    close(fds_config[0]);
    close(fds_data[1]);
}

TEST(ShellSubscriberTest, testMaxSubscriptionsGuard) {
    sp<MockUidMap> uidMap = new NaggyMock<MockUidMap>();
    sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
    sp<ShellSubscriber> shellManager = new ShellSubscriber(uidMap, pullerManager);

    // create a simple config to get screen events
    ShellSubscription config;
    config.add_pushed()->set_atom_id(SCREEN_STATE_CHANGED);

    size_t bufferSize = config.ByteSize();
    vector<uint8_t> buffer(bufferSize);
    config.SerializeToArray(&buffer[0], bufferSize);

    size_t maxSubs = shellManager->getMaxSubscriptions();
    int fds_configs[maxSubs + 1][2];
    int fds_datas[maxSubs + 1][2];
    for (int i = 0; i < maxSubs; i++) {
        // set up 2 pipes for read/write config and data
        ASSERT_EQ(0, pipe2(fds_configs[i], O_CLOEXEC));
        ASSERT_EQ(0, pipe2(fds_datas[i], O_CLOEXEC));

        // write the config to pipe, first write size of the config
        write(fds_configs[i][1], &bufferSize, sizeof(bufferSize));
        // then write config itself
        write(fds_configs[i][1], buffer.data(), bufferSize);
        close(fds_configs[i][1]);

        EXPECT_TRUE(shellManager->startNewSubscription(fds_configs[i][0], fds_datas[i][1],
                                                       /*timeoutSec=*/-1));
        close(fds_configs[i][0]);
        close(fds_datas[i][1]);
    }
    ASSERT_EQ(0, pipe2(fds_configs[maxSubs], O_CLOEXEC));
    ASSERT_EQ(0, pipe2(fds_datas[maxSubs], O_CLOEXEC));

    // write the config to pipe, first write size of the config
    write(fds_configs[maxSubs][1], &bufferSize, sizeof(bufferSize));
    // then write config itself
    write(fds_configs[maxSubs][1], buffer.data(), bufferSize);
    close(fds_configs[maxSubs][1]);

    EXPECT_FALSE(shellManager->startNewSubscription(fds_configs[maxSubs][0], fds_datas[maxSubs][1],
                                                    /*timeoutSec=*/-1));
    close(fds_configs[maxSubs][0]);
    close(fds_datas[maxSubs][1]);

    // Not closing fds_datas[i][0] because this causes writes within ShellSubscriberClient to hang
}

TEST(ShellSubscriberTest, testDifferentConfigs) {
    sp<MockUidMap> uidMap = new NaggyMock<MockUidMap>();
    sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();
    sp<ShellSubscriber> shellManager = new ShellSubscriber(uidMap, pullerManager);

    // number of different configs
    int numConfigs = 2;

    // create a simple config to get screen events
    ShellSubscription configs[numConfigs];
    configs[0].add_pushed()->set_atom_id(SCREEN_STATE_CHANGED);
    configs[1].add_pushed()->set_atom_id(PLUGGED_STATE_CHANGED);

    vector<vector<uint8_t>> configBuffers;
    for (int i = 0; i < numConfigs; i++) {
        size_t bufferSize = configs[i].ByteSize();
        vector<uint8_t> buffer(bufferSize);
        configs[i].SerializeToArray(&buffer[0], bufferSize);
        configBuffers.push_back(buffer);
    }

    int fds_configs[numConfigs][2];
    int fds_datas[numConfigs][2];
    for (int i = 0; i < numConfigs; i++) {
        // set up 2 pipes for read/write config and data
        ASSERT_EQ(0, pipe2(fds_configs[i], O_CLOEXEC));
        ASSERT_EQ(0, pipe2(fds_datas[i], O_CLOEXEC));

        size_t configSize = configBuffers[i].size();
        // write the config to pipe, first write size of the config
        write(fds_configs[i][1], &configSize, sizeof(configSize));
        // then write config itself
        write(fds_configs[i][1], configBuffers[i].data(), configSize);
        close(fds_configs[i][1]);

        EXPECT_TRUE(shellManager->startNewSubscription(fds_configs[i][0], fds_datas[i][1],
                                                       /*timeoutSec=*/-1));
        close(fds_configs[i][0]);
        close(fds_datas[i][1]);
    }

    // send a log event that matches the config.
    vector<std::shared_ptr<LogEvent>> pushedList = getPushedEvents();
    for (const auto& event : pushedList) {
        shellManager->onLogEvent(*event);
    }

    // Validate Config 1
    ShellData actual1 = readData(fds_datas[0][0]);
    ShellData expected1;
    expected1.add_atom()->mutable_screen_state_changed()->set_state(
            ::android::view::DisplayStateEnum::DISPLAY_STATE_ON);
    expected1.add_elapsed_timestamp_nanos(pushedList[0]->GetElapsedTimestampNs());
    EXPECT_THAT(expected1, EqShellData(actual1));

    ShellData actual2 = readData(fds_datas[0][0]);
    ShellData expected2;
    expected2.add_atom()->mutable_screen_state_changed()->set_state(
            ::android::view::DisplayStateEnum::DISPLAY_STATE_OFF);
    expected2.add_elapsed_timestamp_nanos(pushedList[1]->GetElapsedTimestampNs());
    EXPECT_THAT(expected2, EqShellData(actual2));

    // Validate Config 2, repeating the process
    ShellData actual3 = readData(fds_datas[1][0]);
    ShellData expected3;
    expected3.add_atom()->mutable_plugged_state_changed()->set_state(
            BatteryPluggedStateEnum::BATTERY_PLUGGED_USB);
    expected3.add_elapsed_timestamp_nanos(pushedList[2]->GetElapsedTimestampNs());
    EXPECT_THAT(expected3, EqShellData(actual3));

    ShellData actual4 = readData(fds_datas[1][0]);
    ShellData expected4;
    expected4.add_atom()->mutable_plugged_state_changed()->set_state(
            BatteryPluggedStateEnum::BATTERY_PLUGGED_NONE);
    expected4.add_elapsed_timestamp_nanos(pushedList[3]->GetElapsedTimestampNs());
    EXPECT_THAT(expected4, EqShellData(actual4));

    // Not closing fds_datas[i][0] because this causes writes within ShellSubscriberClient to hang
}

TEST(ShellSubscriberTest, testPushedSubscriptionRestrictedEvent) {
    sp<MockUidMap> uidMap = new NaggyMock<MockUidMap>();
    sp<MockStatsPullerManager> pullerManager = new StrictMock<MockStatsPullerManager>();

    std::vector<shared_ptr<LogEvent>> pushedList;
    pushedList.push_back(CreateRestrictedLogEvent(/*atomTag=*/10, /*timestamp=*/1000));

    // create a simple config to get screen events
    ShellSubscription config;
    config.add_pushed()->set_atom_id(10);

    // expect empty data
    vector<ShellData> expectedData;

    // Test with single client
    TRACE_CALL(runShellTest, config, uidMap, pullerManager, pushedList, expectedData,
               kSingleClient);

    // Test with multiple client
    TRACE_CALL(runShellTest, config, uidMap, pullerManager, pushedList, expectedData, kNumClients);
}

#else
GTEST_LOG_(INFO) << "This test does nothing.\n";
#endif

}  // namespace statsd
}  // namespace os
}  // namespace android