diff --git a/cmds/statsd/Android.bp b/cmds/statsd/Android.bp
index cb27325..8af925a 100644
--- a/cmds/statsd/Android.bp
+++ b/cmds/statsd/Android.bp
@@ -230,6 +230,7 @@
         "tests/e2e/Anomaly_duration_sum_e2e_test.cpp",
         "tests/e2e/Attribution_e2e_test.cpp",
         "tests/e2e/ConfigTtl_e2e_test.cpp",
+        "tests/e2e/CountMetric_e2e_test.cpp",
         "tests/e2e/DurationMetric_e2e_test.cpp",
         "tests/e2e/GaugeMetric_e2e_pull_test.cpp",
         "tests/e2e/GaugeMetric_e2e_push_test.cpp",
diff --git a/cmds/statsd/src/StatsLogProcessor.cpp b/cmds/statsd/src/StatsLogProcessor.cpp
index ff7416c..6c3dff2 100644
--- a/cmds/statsd/src/StatsLogProcessor.cpp
+++ b/cmds/statsd/src/StatsLogProcessor.cpp
@@ -16,24 +16,26 @@
 
 #define DEBUG false  // STOPSHIP if true
 #include "Log.h"
-#include "statslog.h"
+
+#include "StatsLogProcessor.h"
 
 #include <android-base/file.h>
 #include <dirent.h>
 #include <frameworks/base/cmds/statsd/src/active_config_list.pb.h>
-#include "StatsLogProcessor.h"
+#include <log/log_event_list.h>
+#include <utils/Errors.h>
+#include <utils/SystemClock.h>
+
 #include "android-base/stringprintf.h"
 #include "external/StatsPullerManager.h"
 #include "guardrail/StatsdStats.h"
 #include "metrics/CountMetricProducer.h"
+#include "state/StateManager.h"
 #include "stats_log_util.h"
 #include "stats_util.h"
+#include "statslog.h"
 #include "storage/StorageManager.h"
 
-#include <log/log_event_list.h>
-#include <utils/Errors.h>
-#include <utils/SystemClock.h>
-
 using namespace android;
 using android::base::StringPrintf;
 using android::util::FIELD_COUNT_REPEATED;
@@ -218,6 +220,8 @@
         onIsolatedUidChangedEventLocked(*event);
     }
 
+    StateManager::getInstance().onLogEvent(*event);
+
     if (mMetricsManagers.empty()) {
         return;
     }
diff --git a/cmds/statsd/src/StatsLogProcessor.h b/cmds/statsd/src/StatsLogProcessor.h
index 3d002d2..8292a3a 100644
--- a/cmds/statsd/src/StatsLogProcessor.h
+++ b/cmds/statsd/src/StatsLogProcessor.h
@@ -262,6 +262,10 @@
     FRIEND_TEST(MetricActivationE2eTest, TestCountMetricWithSameDeactivation);
     FRIEND_TEST(MetricActivationE2eTest, TestCountMetricWithTwoMetricsTwoDeactivations);
 
+    FRIEND_TEST(CountMetricE2eTest, TestWithSimpleState);
+    FRIEND_TEST(CountMetricE2eTest, TestWithMappedState);
+    FRIEND_TEST(CountMetricE2eTest, TestWithMultipleStates);
+
     FRIEND_TEST(DurationMetricE2eTest, TestOneBucket);
     FRIEND_TEST(DurationMetricE2eTest, TestTwoBuckets);
     FRIEND_TEST(DurationMetricE2eTest, TestWithActivation);
diff --git a/cmds/statsd/src/metrics/CountMetricProducer.cpp b/cmds/statsd/src/metrics/CountMetricProducer.cpp
index b5c8e35..4a06387 100644
--- a/cmds/statsd/src/metrics/CountMetricProducer.cpp
+++ b/cmds/statsd/src/metrics/CountMetricProducer.cpp
@@ -18,13 +18,15 @@
 #include "Log.h"
 
 #include "CountMetricProducer.h"
-#include "guardrail/StatsdStats.h"
-#include "stats_util.h"
-#include "stats_log_util.h"
 
+#include <inttypes.h>
 #include <limits.h>
 #include <stdlib.h>
 
+#include "guardrail/StatsdStats.h"
+#include "stats_log_util.h"
+#include "stats_util.h"
+
 using android::util::FIELD_COUNT_REPEATED;
 using android::util::FIELD_TYPE_BOOL;
 using android::util::FIELD_TYPE_FLOAT;
@@ -65,16 +67,16 @@
 const int FIELD_ID_START_BUCKET_ELAPSED_MILLIS = 5;
 const int FIELD_ID_END_BUCKET_ELAPSED_MILLIS = 6;
 
-CountMetricProducer::CountMetricProducer(const ConfigKey& key, const CountMetric& metric,
-                                         const int conditionIndex,
-                                         const sp<ConditionWizard>& wizard,
-                                         const int64_t timeBaseNs, const int64_t startTimeNs,
-                                         const unordered_map<int, shared_ptr<Activation>>&
-                                                 eventActivationMap,
-                                         const unordered_map<int, vector<shared_ptr<Activation>>>&
-                                                 eventDeactivationMap)
+CountMetricProducer::CountMetricProducer(
+        const ConfigKey& key, const CountMetric& metric, const int conditionIndex,
+        const sp<ConditionWizard>& wizard, const int64_t timeBaseNs, const int64_t startTimeNs,
+
+        const unordered_map<int, shared_ptr<Activation>>& eventActivationMap,
+        const unordered_map<int, vector<shared_ptr<Activation>>>& eventDeactivationMap,
+        const vector<int>& slicedStateAtoms,
+        const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap)
     : MetricProducer(metric.id(), key, timeBaseNs, conditionIndex, wizard, eventActivationMap,
-                     eventDeactivationMap) {
+                     eventDeactivationMap, slicedStateAtoms, stateGroupMap) {
     if (metric.has_bucket()) {
         mBucketSizeNs =
                 TimeUnitToBucketSizeInMillisGuardrailed(key.GetUid(), metric.bucket()) * 1000000;
@@ -100,6 +102,8 @@
         mConditionSliced = true;
     }
 
+    // TODO(tsaichristine): b/142124705 handle metric state links
+
     flushIfNeededLocked(startTimeNs);
     // Adjust start for partial bucket
     mCurrentBucketStartTimeNs = startTimeNs;
@@ -112,6 +116,12 @@
     VLOG("~CountMetricProducer() called");
 }
 
+void CountMetricProducer::onStateChanged(int atomId, const HashableDimensionKey& primaryKey,
+                                         int oldState, int newState) {
+    VLOG("CountMetric %lld onStateChanged State%d, key %s, %d -> %d", (long long)mMetricId, atomId,
+         primaryKey.toString().c_str(), oldState, newState);
+}
+
 void CountMetricProducer::dumpStatesLocked(FILE* out, bool verbose) const {
     if (mCurrentSlicedCounter == nullptr ||
         mCurrentSlicedCounter->size() == 0) {
diff --git a/cmds/statsd/src/metrics/CountMetricProducer.h b/cmds/statsd/src/metrics/CountMetricProducer.h
index 1cdc976..61e0892 100644
--- a/cmds/statsd/src/metrics/CountMetricProducer.h
+++ b/cmds/statsd/src/metrics/CountMetricProducer.h
@@ -17,15 +17,16 @@
 #ifndef COUNT_METRIC_PRODUCER_H
 #define COUNT_METRIC_PRODUCER_H
 
-#include <unordered_map>
-
 #include <android/util/ProtoOutputStream.h>
 #include <gtest/gtest_prod.h>
-#include "../anomaly/AnomalyTracker.h"
-#include "../condition/ConditionTracker.h"
-#include "../matchers/matcher_util.h"
+
+#include <unordered_map>
+
 #include "MetricProducer.h"
+#include "anomaly/AnomalyTracker.h"
+#include "condition/ConditionTracker.h"
 #include "frameworks/base/cmds/statsd/src/statsd_config.pb.h"
+#include "matchers/matcher_util.h"
 #include "stats_util.h"
 
 namespace android {
@@ -40,16 +41,20 @@
 
 class CountMetricProducer : public MetricProducer {
 public:
-    CountMetricProducer(const ConfigKey& key, const CountMetric& countMetric,
-                        const int conditionIndex, const sp<ConditionWizard>& wizard,
-                        const int64_t timeBaseNs, const int64_t startTimeNs,
-                        const std::unordered_map<int, std::shared_ptr<Activation>>&
-                                eventActivationMap = {},
-                        const std::unordered_map<int, std::vector<std::shared_ptr<Activation>>>&
-                                eventDeactivationMap = {});
+    CountMetricProducer(
+            const ConfigKey& key, const CountMetric& countMetric, const int conditionIndex,
+            const sp<ConditionWizard>& wizard, const int64_t timeBaseNs, const int64_t startTimeNs,
+            const std::unordered_map<int, std::shared_ptr<Activation>>& eventActivationMap = {},
+            const std::unordered_map<int, std::vector<std::shared_ptr<Activation>>>&
+                    eventDeactivationMap = {},
+            const vector<int>& slicedStateAtoms = {},
+            const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap = {});
 
     virtual ~CountMetricProducer();
 
+    void onStateChanged(int atomId, const HashableDimensionKey& primaryKey, int oldState,
+                        int newState) override;
+
 protected:
     void onMatchedLogEventInternalLocked(
             const size_t matcherIndex, const MetricDimensionKey& eventKey,
diff --git a/cmds/statsd/src/metrics/DurationMetricProducer.cpp b/cmds/statsd/src/metrics/DurationMetricProducer.cpp
index 31b90f3..ab2a1c3 100644
--- a/cmds/statsd/src/metrics/DurationMetricProducer.cpp
+++ b/cmds/statsd/src/metrics/DurationMetricProducer.cpp
@@ -69,9 +69,11 @@
         const bool nesting, const sp<ConditionWizard>& wizard,
         const FieldMatcher& internalDimensions, const int64_t timeBaseNs, const int64_t startTimeNs,
         const unordered_map<int, shared_ptr<Activation>>& eventActivationMap,
-        const unordered_map<int, vector<shared_ptr<Activation>>>& eventDeactivationMap)
+        const unordered_map<int, vector<shared_ptr<Activation>>>& eventDeactivationMap,
+        const vector<int>& slicedStateAtoms,
+        const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap)
     : MetricProducer(metric.id(), key, timeBaseNs, conditionIndex, wizard, eventActivationMap,
-                     eventDeactivationMap),
+                     eventDeactivationMap, slicedStateAtoms, stateGroupMap),
       mAggregationType(metric.aggregation_type()),
       mStartIndex(startIndex),
       mStopIndex(stopIndex),
diff --git a/cmds/statsd/src/metrics/DurationMetricProducer.h b/cmds/statsd/src/metrics/DurationMetricProducer.h
index 0592b18..7457d7f 100644
--- a/cmds/statsd/src/metrics/DurationMetricProducer.h
+++ b/cmds/statsd/src/metrics/DurationMetricProducer.h
@@ -38,16 +38,16 @@
 
 class DurationMetricProducer : public MetricProducer {
 public:
-    DurationMetricProducer(const ConfigKey& key, const DurationMetric& durationMetric,
-                           const int conditionIndex, const size_t startIndex,
-                           const size_t stopIndex, const size_t stopAllIndex, const bool nesting,
-                           const sp<ConditionWizard>& wizard,
-                           const FieldMatcher& internalDimensions, const int64_t timeBaseNs,
-                           const int64_t startTimeNs,
-                           const unordered_map<int, shared_ptr<Activation>>&
-                                   eventActivationMap = {},
-                           const unordered_map<int, vector<shared_ptr<Activation>>>&
-                                   eventDeactivationMap = {});
+    DurationMetricProducer(
+            const ConfigKey& key, const DurationMetric& durationMetric, const int conditionIndex,
+            const size_t startIndex, const size_t stopIndex, const size_t stopAllIndex,
+            const bool nesting, const sp<ConditionWizard>& wizard,
+            const FieldMatcher& internalDimensions, const int64_t timeBaseNs,
+            const int64_t startTimeNs,
+            const unordered_map<int, shared_ptr<Activation>>& eventActivationMap = {},
+            const unordered_map<int, vector<shared_ptr<Activation>>>& eventDeactivationMap = {},
+            const vector<int>& slicedStateAtoms = {},
+            const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap = {});
 
     virtual ~DurationMetricProducer();
 
diff --git a/cmds/statsd/src/metrics/EventMetricProducer.cpp b/cmds/statsd/src/metrics/EventMetricProducer.cpp
index a60a916..32eb077 100644
--- a/cmds/statsd/src/metrics/EventMetricProducer.cpp
+++ b/cmds/statsd/src/metrics/EventMetricProducer.cpp
@@ -52,16 +52,15 @@
 const int FIELD_ID_ELAPSED_TIMESTAMP_NANOS = 1;
 const int FIELD_ID_ATOMS = 2;
 
-EventMetricProducer::EventMetricProducer(const ConfigKey& key, const EventMetric& metric,
-                                         const int conditionIndex,
-                                         const sp<ConditionWizard>& wizard,
-                                         const int64_t startTimeNs,
-                                         const unordered_map<int, shared_ptr<Activation>>&
-                                                 eventActivationMap,
-                                         const unordered_map<int, vector<shared_ptr<Activation>>>&
-                                                 eventDeactivationMap)
+EventMetricProducer::EventMetricProducer(
+        const ConfigKey& key, const EventMetric& metric, const int conditionIndex,
+        const sp<ConditionWizard>& wizard, const int64_t startTimeNs,
+        const unordered_map<int, shared_ptr<Activation>>& eventActivationMap,
+        const unordered_map<int, vector<shared_ptr<Activation>>>& eventDeactivationMap,
+        const vector<int>& slicedStateAtoms,
+        const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap)
     : MetricProducer(metric.id(), key, startTimeNs, conditionIndex, wizard, eventActivationMap,
-                     eventDeactivationMap) {
+                     eventDeactivationMap, slicedStateAtoms, stateGroupMap) {
     if (metric.links().size() > 0) {
         for (const auto& link : metric.links()) {
             Metric2Condition mc;
diff --git a/cmds/statsd/src/metrics/EventMetricProducer.h b/cmds/statsd/src/metrics/EventMetricProducer.h
index aab53c8..dca37e8 100644
--- a/cmds/statsd/src/metrics/EventMetricProducer.h
+++ b/cmds/statsd/src/metrics/EventMetricProducer.h
@@ -33,13 +33,14 @@
 
 class EventMetricProducer : public MetricProducer {
 public:
-    EventMetricProducer(const ConfigKey& key, const EventMetric& eventMetric,
-                        const int conditionIndex, const sp<ConditionWizard>& wizard,
-                        const int64_t startTimeNs,
-                        const std::unordered_map<int, std::shared_ptr<Activation>>&
-                                eventActivationMap = {},
-                        const std::unordered_map<int, std::vector<std::shared_ptr<Activation>>>&
-                                eventDeactivationMap = {});
+    EventMetricProducer(
+            const ConfigKey& key, const EventMetric& eventMetric, const int conditionIndex,
+            const sp<ConditionWizard>& wizard, const int64_t startTimeNs,
+            const std::unordered_map<int, std::shared_ptr<Activation>>& eventActivationMap = {},
+            const std::unordered_map<int, std::vector<std::shared_ptr<Activation>>>&
+                    eventDeactivationMap = {},
+            const vector<int>& slicedStateAtoms = {},
+            const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap = {});
 
     virtual ~EventMetricProducer();
 
diff --git a/cmds/statsd/src/metrics/GaugeMetricProducer.cpp b/cmds/statsd/src/metrics/GaugeMetricProducer.cpp
index e409b6fb..d0f88a8 100644
--- a/cmds/statsd/src/metrics/GaugeMetricProducer.cpp
+++ b/cmds/statsd/src/metrics/GaugeMetricProducer.cpp
@@ -74,9 +74,11 @@
         const int atomId, const int64_t timeBaseNs, const int64_t startTimeNs,
         const sp<StatsPullerManager>& pullerManager,
         const unordered_map<int, shared_ptr<Activation>>& eventActivationMap,
-        const unordered_map<int, vector<shared_ptr<Activation>>>& eventDeactivationMap)
+        const unordered_map<int, vector<shared_ptr<Activation>>>& eventDeactivationMap,
+        const vector<int>& slicedStateAtoms,
+        const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap)
     : MetricProducer(metric.id(), key, timeBaseNs, conditionIndex, wizard, eventActivationMap,
-            eventDeactivationMap),
+                     eventDeactivationMap, slicedStateAtoms, stateGroupMap),
       mWhatMatcherIndex(whatMatcherIndex),
       mEventMatcherWizard(matcherWizard),
       mPullerManager(pullerManager),
diff --git a/cmds/statsd/src/metrics/GaugeMetricProducer.h b/cmds/statsd/src/metrics/GaugeMetricProducer.h
index dfe1d56..640a02a 100644
--- a/cmds/statsd/src/metrics/GaugeMetricProducer.h
+++ b/cmds/statsd/src/metrics/GaugeMetricProducer.h
@@ -56,16 +56,17 @@
 // producer always reports the guage at the earliest time of the bucket when the condition is met.
 class GaugeMetricProducer : public virtual MetricProducer, public virtual PullDataReceiver {
 public:
-    GaugeMetricProducer(const ConfigKey& key, const GaugeMetric& gaugeMetric,
-                        const int conditionIndex, const sp<ConditionWizard>& conditionWizard,
-                        const int whatMatcherIndex,const sp<EventMatcherWizard>& matcherWizard,
-                        const int pullTagId, const int triggerAtomId, const int atomId,
-                        const int64_t timeBaseNs, const int64_t startTimeNs,
-                        const sp<StatsPullerManager>& pullerManager,
-                        const std::unordered_map<int, std::shared_ptr<Activation>>&
-                                eventActivationMap = {},
-                        const std::unordered_map<int, std::vector<std::shared_ptr<Activation>>>&
-                                eventDeactivationMap = {});
+    GaugeMetricProducer(
+            const ConfigKey& key, const GaugeMetric& gaugeMetric, const int conditionIndex,
+            const sp<ConditionWizard>& conditionWizard, const int whatMatcherIndex,
+            const sp<EventMatcherWizard>& matcherWizard, const int pullTagId,
+            const int triggerAtomId, const int atomId, const int64_t timeBaseNs,
+            const int64_t startTimeNs, const sp<StatsPullerManager>& pullerManager,
+            const std::unordered_map<int, std::shared_ptr<Activation>>& eventActivationMap = {},
+            const std::unordered_map<int, std::vector<std::shared_ptr<Activation>>>&
+                    eventDeactivationMap = {},
+            const vector<int>& slicedStateAtoms = {},
+            const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap = {});
 
     virtual ~GaugeMetricProducer();
 
diff --git a/cmds/statsd/src/metrics/MetricProducer.cpp b/cmds/statsd/src/metrics/MetricProducer.cpp
index 3426a19..2a700ef 100644
--- a/cmds/statsd/src/metrics/MetricProducer.cpp
+++ b/cmds/statsd/src/metrics/MetricProducer.cpp
@@ -45,23 +45,27 @@
         const int conditionIndex, const sp<ConditionWizard>& wizard,
         const std::unordered_map<int, std::shared_ptr<Activation>>& eventActivationMap,
         const std::unordered_map<int, std::vector<std::shared_ptr<Activation>>>&
-                eventDeactivationMap)
-        : mMetricId(metricId),
-          mConfigKey(key),
-          mTimeBaseNs(timeBaseNs),
-          mCurrentBucketStartTimeNs(timeBaseNs),
-          mCurrentBucketNum(0),
-          mCondition(initialCondition(conditionIndex)),
-          mConditionTrackerIndex(conditionIndex),
-          mConditionSliced(false),
-          mWizard(wizard),
-          mContainANYPositionInDimensionsInWhat(false),
-          mSliceByPositionALL(false),
-          mHasLinksToAllConditionDimensionsInTracker(false),
-          mEventActivationMap(eventActivationMap),
-          mEventDeactivationMap(eventDeactivationMap),
-          mIsActive(mEventActivationMap.empty()) {
-    }
+                eventDeactivationMap,
+        const vector<int>& slicedStateAtoms,
+        const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap)
+    : mMetricId(metricId),
+      mConfigKey(key),
+      mTimeBaseNs(timeBaseNs),
+      mCurrentBucketStartTimeNs(timeBaseNs),
+      mCurrentBucketNum(0),
+      mCondition(initialCondition(conditionIndex)),
+      mConditionTrackerIndex(conditionIndex),
+      mConditionSliced(false),
+      mWizard(wizard),
+      mContainANYPositionInDimensionsInWhat(false),
+      mSliceByPositionALL(false),
+      mHasLinksToAllConditionDimensionsInTracker(false),
+      mEventActivationMap(eventActivationMap),
+      mEventDeactivationMap(eventDeactivationMap),
+      mIsActive(mEventActivationMap.empty()),
+      mSlicedStateAtoms(slicedStateAtoms),
+      mStateGroupMap(stateGroupMap) {
+}
 
 void MetricProducer::onMatchedLogEventLocked(const size_t matcherIndex, const LogEvent& event) {
     if (!mIsActive) {
diff --git a/cmds/statsd/src/metrics/MetricProducer.h b/cmds/statsd/src/metrics/MetricProducer.h
index 1e1eb69..a72de22 100644
--- a/cmds/statsd/src/metrics/MetricProducer.h
+++ b/cmds/statsd/src/metrics/MetricProducer.h
@@ -17,19 +17,19 @@
 #ifndef METRIC_PRODUCER_H
 #define METRIC_PRODUCER_H
 
-#include <shared_mutex>
-
 #include <frameworks/base/cmds/statsd/src/active_config_list.pb.h>
+#include <log/logprint.h>
+#include <utils/RefBase.h>
+
+#include <unordered_map>
+
 #include "HashableDimensionKey.h"
 #include "anomaly/AnomalyTracker.h"
 #include "condition/ConditionWizard.h"
 #include "config/ConfigKey.h"
 #include "matchers/matcher_util.h"
 #include "packages/PackageInfoListener.h"
-
-#include <log/logprint.h>
-#include <utils/RefBase.h>
-#include <unordered_map>
+#include "state/StateListener.h"
 
 namespace android {
 namespace os {
@@ -86,13 +86,15 @@
 // writing the report to dropbox. MetricProducers should respond to package changes as required in
 // PackageInfoListener, but if none of the metrics are slicing by package name, then the update can
 // be a no-op.
-class MetricProducer : public virtual PackageInfoListener {
+class MetricProducer : public virtual PackageInfoListener, public virtual StateListener {
 public:
     MetricProducer(const int64_t& metricId, const ConfigKey& key, const int64_t timeBaseNs,
                    const int conditionIndex, const sp<ConditionWizard>& wizard,
                    const std::unordered_map<int, std::shared_ptr<Activation>>& eventActivationMap,
                    const std::unordered_map<int, std::vector<std::shared_ptr<Activation>>>&
-                           eventDeactivationMap);
+                           eventDeactivationMap,
+                   const vector<int>& slicedStateAtoms,
+                   const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap);
 
     virtual ~MetricProducer(){};
 
@@ -151,6 +153,9 @@
         return mConditionSliced;
     };
 
+    void onStateChanged(int atomId, const HashableDimensionKey& primaryKey, int oldState,
+                        int newState){};
+
     // Output the metrics data to [protoOutput]. All metrics reports end with the same timestamp.
     // This method clears all the past buckets.
     void onDumpReport(const int64_t dumpTimeNs,
@@ -230,6 +235,11 @@
         return mBucketSizeNs;
     }
 
+    inline const std::vector<int> getSlicedStateAtoms() {
+        std::lock_guard<std::mutex> lock(mMutex);
+        return mSlicedStateAtoms;
+    }
+
     /* If alert is valid, adds an AnomalyTracker and returns it. If invalid, returns nullptr. */
     virtual sp<AnomalyTracker> addAnomalyTracker(const Alert &alert,
                                                  const sp<AlarmMonitor>& anomalyAlarmMonitor) {
@@ -381,6 +391,16 @@
 
     bool mIsActive;
 
+    // The slice_by_state atom ids defined in statsd_config.
+    std::vector<int> mSlicedStateAtoms;
+
+    // Maps atom ids and state values to group_ids (<atom_id, <value, group_id>>).
+    std::unordered_map<int, std::unordered_map<int, int64_t>> mStateGroupMap;
+
+    FRIEND_TEST(CountMetricE2eTest, TestWithSimpleState);
+    FRIEND_TEST(CountMetricE2eTest, TestWithMappedState);
+    FRIEND_TEST(CountMetricE2eTest, TestWithMultipleStates);
+
     FRIEND_TEST(DurationMetricE2eTest, TestOneBucket);
     FRIEND_TEST(DurationMetricE2eTest, TestTwoBuckets);
     FRIEND_TEST(DurationMetricE2eTest, TestWithActivation);
diff --git a/cmds/statsd/src/metrics/MetricsManager.cpp b/cmds/statsd/src/metrics/MetricsManager.cpp
index 963205e..7bae4b9 100644
--- a/cmds/statsd/src/metrics/MetricsManager.cpp
+++ b/cmds/statsd/src/metrics/MetricsManager.cpp
@@ -15,8 +15,12 @@
  */
 #define DEBUG false  // STOPSHIP if true
 #include "Log.h"
+
 #include "MetricsManager.h"
-#include "statslog.h"
+
+#include <log/logprint.h>
+#include <private/android_filesystem_config.h>
+#include <utils/SystemClock.h>
 
 #include "CountMetricProducer.h"
 #include "condition/CombinationConditionTracker.h"
@@ -25,12 +29,10 @@
 #include "matchers/CombinationLogMatchingTracker.h"
 #include "matchers/SimpleLogMatchingTracker.h"
 #include "metrics_manager_util.h"
-#include "stats_util.h"
+#include "state/StateManager.h"
 #include "stats_log_util.h"
-
-#include <log/logprint.h>
-#include <private/android_filesystem_config.h>
-#include <utils/SystemClock.h>
+#include "stats_util.h"
+#include "statslog.h"
 
 using android::util::FIELD_COUNT_REPEATED;
 using android::util::FIELD_TYPE_INT32;
@@ -149,6 +151,12 @@
 }
 
 MetricsManager::~MetricsManager() {
+    for (auto it : mAllMetricProducers) {
+        for (int atomId : it->getSlicedStateAtoms()) {
+            StateManager::getInstance().unregisterListener(atomId, it);
+        }
+    }
+
     VLOG("~MetricsManager()");
 }
 
diff --git a/cmds/statsd/src/metrics/MetricsManager.h b/cmds/statsd/src/metrics/MetricsManager.h
index 3dad614..d184121 100644
--- a/cmds/statsd/src/metrics/MetricsManager.h
+++ b/cmds/statsd/src/metrics/MetricsManager.h
@@ -282,6 +282,10 @@
             TestActivationOnBootMultipleActivationsDifferentActivationTypes);
     FRIEND_TEST(StatsLogProcessorTest, TestActivationsPersistAcrossSystemServerRestart);
 
+    FRIEND_TEST(CountMetricE2eTest, TestWithSimpleState);
+    FRIEND_TEST(CountMetricE2eTest, TestWithMappedState);
+    FRIEND_TEST(CountMetricE2eTest, TestWithMultipleStates);
+
     FRIEND_TEST(DurationMetricE2eTest, TestOneBucket);
     FRIEND_TEST(DurationMetricE2eTest, TestTwoBuckets);
     FRIEND_TEST(DurationMetricE2eTest, TestWithActivation);
diff --git a/cmds/statsd/src/metrics/ValueMetricProducer.cpp b/cmds/statsd/src/metrics/ValueMetricProducer.cpp
index 7fe5a83..6fd0327 100644
--- a/cmds/statsd/src/metrics/ValueMetricProducer.cpp
+++ b/cmds/statsd/src/metrics/ValueMetricProducer.cpp
@@ -83,9 +83,11 @@
         const sp<EventMatcherWizard>& matcherWizard, const int pullTagId, const int64_t timeBaseNs,
         const int64_t startTimeNs, const sp<StatsPullerManager>& pullerManager,
         const unordered_map<int, shared_ptr<Activation>>& eventActivationMap,
-        const unordered_map<int, vector<shared_ptr<Activation>>>& eventDeactivationMap)
+        const unordered_map<int, vector<shared_ptr<Activation>>>& eventDeactivationMap,
+        const vector<int>& slicedStateAtoms,
+        const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap)
     : MetricProducer(metric.id(), key, timeBaseNs, conditionIndex, conditionWizard,
-                     eventActivationMap, eventDeactivationMap),
+                     eventActivationMap, eventDeactivationMap, slicedStateAtoms, stateGroupMap),
       mWhatMatcherIndex(whatMatcherIndex),
       mEventMatcherWizard(matcherWizard),
       mPullerManager(pullerManager),
diff --git a/cmds/statsd/src/metrics/ValueMetricProducer.h b/cmds/statsd/src/metrics/ValueMetricProducer.h
index d7cd397..206e602 100644
--- a/cmds/statsd/src/metrics/ValueMetricProducer.h
+++ b/cmds/statsd/src/metrics/ValueMetricProducer.h
@@ -52,15 +52,17 @@
 // - an alarm set to the end of the bucket
 class ValueMetricProducer : public virtual MetricProducer, public virtual PullDataReceiver {
 public:
-    ValueMetricProducer(const ConfigKey& key, const ValueMetric& valueMetric,
-                        const int conditionIndex, const sp<ConditionWizard>& conditionWizard,
-                        const int whatMatcherIndex, const sp<EventMatcherWizard>& matcherWizard,
-                        const int pullTagId, const int64_t timeBaseNs, const int64_t startTimeNs,
-                        const sp<StatsPullerManager>& pullerManager,
-                        const std::unordered_map<int, std::shared_ptr<Activation>>&
-                                eventActivationMap = {},
-                        const std::unordered_map<int, std::vector<std::shared_ptr<Activation>>>&
-                                eventDeactivationMap = {});
+    ValueMetricProducer(
+            const ConfigKey& key, const ValueMetric& valueMetric, const int conditionIndex,
+            const sp<ConditionWizard>& conditionWizard, const int whatMatcherIndex,
+            const sp<EventMatcherWizard>& matcherWizard, const int pullTagId,
+            const int64_t timeBaseNs, const int64_t startTimeNs,
+            const sp<StatsPullerManager>& pullerManager,
+            const std::unordered_map<int, std::shared_ptr<Activation>>& eventActivationMap = {},
+            const std::unordered_map<int, std::vector<std::shared_ptr<Activation>>>&
+                    eventDeactivationMap = {},
+            const vector<int>& slicedStateAtoms = {},
+            const unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap = {});
 
     virtual ~ValueMetricProducer();
 
diff --git a/cmds/statsd/src/metrics/metrics_manager_util.cpp b/cmds/statsd/src/metrics/metrics_manager_util.cpp
index 0fee71e..33e162e 100644
--- a/cmds/statsd/src/metrics/metrics_manager_util.cpp
+++ b/cmds/statsd/src/metrics/metrics_manager_util.cpp
@@ -20,24 +20,24 @@
 #include "metrics_manager_util.h"
 #include "MetricProducer.h"
 
-#include "../condition/CombinationConditionTracker.h"
-#include "../condition/SimpleConditionTracker.h"
-#include "../condition/StateConditionTracker.h"
-#include "../external/StatsPullerManager.h"
-#include "../matchers/CombinationLogMatchingTracker.h"
-#include "../matchers/SimpleLogMatchingTracker.h"
-#include "../matchers/EventMatcherWizard.h"
-#include "../metrics/CountMetricProducer.h"
-#include "../metrics/DurationMetricProducer.h"
-#include "../metrics/EventMetricProducer.h"
-#include "../metrics/GaugeMetricProducer.h"
-#include "../metrics/ValueMetricProducer.h"
+#include <inttypes.h>
 
+#include "condition/CombinationConditionTracker.h"
+#include "condition/SimpleConditionTracker.h"
+#include "condition/StateConditionTracker.h"
+#include "external/StatsPullerManager.h"
+#include "matchers/CombinationLogMatchingTracker.h"
+#include "matchers/EventMatcherWizard.h"
+#include "matchers/SimpleLogMatchingTracker.h"
+#include "metrics/CountMetricProducer.h"
+#include "metrics/DurationMetricProducer.h"
+#include "metrics/EventMetricProducer.h"
+#include "metrics/GaugeMetricProducer.h"
+#include "metrics/ValueMetricProducer.h"
+#include "state/StateManager.h"
 #include "stats_util.h"
 #include "statslog.h"
 
-#include <inttypes.h>
-
 using std::set;
 using std::string;
 using std::unordered_map;
@@ -138,6 +138,41 @@
     return true;
 }
 
+// Initializes state data structures for a metric.
+// input:
+// [config]: the input config
+// [stateIds]: the slice_by_state ids for this metric
+// [stateAtomIdMap]: this map contains the mapping from all state ids to atom ids
+// [allStateGroupMaps]: this map contains the mapping from state ids and state
+//                      values to state group ids for all states
+// output:
+// [slicedStateAtoms]: a vector of atom ids of all the slice_by_states
+// [stateGroupMap]: this map should contain the mapping from states ids and state
+//                      values to state group ids for all states that this metric
+//                      is interested in
+bool handleMetricWithStates(
+        const StatsdConfig& config, const ::google::protobuf::RepeatedField<int64_t>& stateIds,
+        const unordered_map<int64_t, int>& stateAtomIdMap,
+        const unordered_map<int64_t, unordered_map<int, int64_t>>& allStateGroupMaps,
+        vector<int>& slicedStateAtoms,
+        unordered_map<int, unordered_map<int, int64_t>>& stateGroupMap) {
+    for (const auto& stateId : stateIds) {
+        auto it = stateAtomIdMap.find(stateId);
+        if (it == stateAtomIdMap.end()) {
+            ALOGW("cannot find State %" PRId64 " in the config", stateId);
+            return false;
+        }
+        int atomId = it->second;
+        slicedStateAtoms.push_back(atomId);
+
+        auto stateIt = allStateGroupMaps.find(stateId);
+        if (stateIt != allStateGroupMaps.end()) {
+            stateGroupMap[atomId] = stateIt->second;
+        }
+    }
+    return true;
+}
+
 // Validates a metricActivation and populates state.
 // EventActivationMap and EventDeactivationMap are supplied to a MetricProducer
 //      to provide the producer with state about its activators and deactivators.
@@ -342,12 +377,32 @@
     return true;
 }
 
+bool initStates(const StatsdConfig& config, unordered_map<int64_t, int>& stateAtomIdMap,
+                unordered_map<int64_t, unordered_map<int, int64_t>>& allStateGroupMaps) {
+    for (int i = 0; i < config.state_size(); i++) {
+        const State& state = config.state(i);
+        const int64_t stateId = state.id();
+        stateAtomIdMap[stateId] = state.atom_id();
+
+        const StateMap& stateMap = state.map();
+        for (auto group : stateMap.group()) {
+            for (auto value : group.value()) {
+                allStateGroupMaps[stateId][value] = group.group_id();
+            }
+        }
+    }
+
+    return true;
+}
+
 bool initMetrics(const ConfigKey& key, const StatsdConfig& config, const int64_t timeBaseTimeNs,
                  const int64_t currentTimeNs, UidMap& uidMap,
                  const sp<StatsPullerManager>& pullerManager,
                  const unordered_map<int64_t, int>& logTrackerMap,
                  const unordered_map<int64_t, int>& conditionTrackerMap,
                  const vector<sp<LogMatchingTracker>>& allAtomMatchers,
+                 const unordered_map<int64_t, int>& stateAtomIdMap,
+                 const unordered_map<int64_t, unordered_map<int, int64_t>>& allStateGroupMaps,
                  vector<sp<ConditionTracker>>& allConditionTrackers,
                  vector<sp<MetricProducer>>& allMetricProducers,
                  unordered_map<int, vector<int>>& conditionToMetricMap,
@@ -398,10 +453,9 @@
 
         int conditionIndex = -1;
         if (metric.has_condition()) {
-            bool good = handleMetricWithConditions(
-                    metric.condition(), metricIndex, conditionTrackerMap, metric.links(),
-                    allConditionTrackers, conditionIndex, conditionToMetricMap);
-            if (!good) {
+            if (!handleMetricWithConditions(metric.condition(), metricIndex, conditionTrackerMap,
+                                            metric.links(), allConditionTrackers, conditionIndex,
+                                            conditionToMetricMap)) {
                 return false;
             }
         } else {
@@ -411,6 +465,18 @@
             }
         }
 
+        std::vector<int> slicedStateAtoms;
+        unordered_map<int, unordered_map<int, int64_t>> stateGroupMap;
+        if (metric.slice_by_state_size() > 0) {
+            if (!handleMetricWithStates(config, metric.slice_by_state(), stateAtomIdMap,
+                                        allStateGroupMaps, slicedStateAtoms, stateGroupMap)) {
+                return false;
+            }
+        }
+
+        // TODO(tsaichristine): add check for unequal number of MetricStateLinks
+        // and slice_by_states
+
         unordered_map<int, shared_ptr<Activation>> eventActivationMap;
         unordered_map<int, vector<shared_ptr<Activation>>> eventDeactivationMap;
         bool success = handleMetricActivation(config, metric.id(), metricIndex,
@@ -421,7 +487,7 @@
 
         sp<MetricProducer> countProducer = new CountMetricProducer(
                 key, metric, conditionIndex, wizard, timeBaseTimeNs, currentTimeNs,
-                eventActivationMap, eventDeactivationMap);
+                eventActivationMap, eventDeactivationMap, slicedStateAtoms, stateGroupMap);
         allMetricProducers.push_back(countProducer);
     }
 
@@ -721,6 +787,13 @@
     }
     for (const auto& it : allMetricProducers) {
         uidMap.addListener(it);
+
+        // Register metrics to StateTrackers
+        for (int atomId : it->getSlicedStateAtoms()) {
+            if (!StateManager::getInstance().registerListener(atomId, it)) {
+                return false;
+            }
+        }
     }
     return true;
 }
@@ -846,6 +919,8 @@
     unordered_map<int64_t, int> logTrackerMap;
     unordered_map<int64_t, int> conditionTrackerMap;
     unordered_map<int64_t, int> metricProducerMap;
+    unordered_map<int64_t, int> stateAtomIdMap;
+    unordered_map<int64_t, unordered_map<int, int64_t>> allStateGroupMaps;
 
     if (!initLogTrackers(config, uidMap, logTrackerMap, allAtomMatchers, allTagIds)) {
         ALOGE("initLogMatchingTrackers failed");
@@ -858,9 +933,13 @@
         ALOGE("initConditionTrackers failed");
         return false;
     }
-
+    if (!initStates(config, stateAtomIdMap, allStateGroupMaps)) {
+        ALOGE("initStates failed");
+        return false;
+    }
     if (!initMetrics(key, config, timeBaseNs, currentTimeNs, uidMap, pullerManager, logTrackerMap,
-                     conditionTrackerMap, allAtomMatchers, allConditionTrackers, allMetricProducers,
+                     conditionTrackerMap, allAtomMatchers, stateAtomIdMap, allStateGroupMaps,
+                     allConditionTrackers, allMetricProducers,
                      conditionToMetricMap, trackerToMetricMap, metricProducerMap,
                      noReportMetricIds, activationAtomTrackerToMetricMap,
                      deactivationAtomTrackerToMetricMap, metricsWithActivation)) {
diff --git a/cmds/statsd/src/metrics/metrics_manager_util.h b/cmds/statsd/src/metrics/metrics_manager_util.h
index 3802948..95b2ab8 100644
--- a/cmds/statsd/src/metrics/metrics_manager_util.h
+++ b/cmds/statsd/src/metrics/metrics_manager_util.h
@@ -68,6 +68,17 @@
                     std::unordered_map<int, std::vector<int>>& trackerToConditionMap,
                     std::unordered_map<int, std::vector<MetricConditionLink>>& eventConditionLinks);
 
+// Initialize State maps using State protos in the config. These maps will
+// eventually be passed to MetricProducers to initialize their state info.
+// input:
+// [config]: the input config
+// output:
+// [stateAtomIdMap]: this map should contain the mapping from state ids to atom ids
+// [allStateGroupMaps]: this map should contain the mapping from states ids and state
+//                      values to state group ids for all states
+bool initStates(const StatsdConfig& config, unordered_map<int64_t, int>& stateAtomIdMap,
+                unordered_map<int64_t, unordered_map<int, int64_t>>& allStateGroupMaps);
+
 // Initialize MetricProducers.
 // input:
 // [key]: the config key that this config belongs to
@@ -75,6 +86,9 @@
 // [timeBaseSec]: start time base for all metrics
 // [logTrackerMap]: LogMatchingTracker name to index mapping from previous step.
 // [conditionTrackerMap]: condition name to index mapping
+// [stateAtomIdMap]: contains the mapping from state ids to atom ids
+// [allStateGroupMaps]: contains the mapping from atom ids and state values to
+//                      state group ids for all states
 // output:
 // [allMetricProducers]: contains the list of sp to the MetricProducers created.
 // [conditionToMetricMap]: contains the mapping from condition tracker index to
@@ -87,6 +101,8 @@
         const std::unordered_map<int64_t, int>& conditionTrackerMap,
         const std::unordered_map<int, std::vector<MetricConditionLink>>& eventConditionLinks,
         const vector<sp<LogMatchingTracker>>& allAtomMatchers,
+        const unordered_map<int64_t, int>& stateAtomIdMap,
+        const unordered_map<int64_t, unordered_map<int, int64_t>>& allStateGroupMaps,
         vector<sp<ConditionTracker>>& allConditionTrackers,
         std::vector<sp<MetricProducer>>& allMetricProducers,
         std::unordered_map<int, std::vector<int>>& conditionToMetricMap,
diff --git a/cmds/statsd/src/state/StateManager.cpp b/cmds/statsd/src/state/StateManager.cpp
index a3059c5..95b2c76 100644
--- a/cmds/statsd/src/state/StateManager.cpp
+++ b/cmds/statsd/src/state/StateManager.cpp
@@ -35,26 +35,25 @@
     }
 }
 
-bool StateManager::registerListener(int stateAtomId, wp<StateListener> listener) {
+bool StateManager::registerListener(int atomId, wp<StateListener> listener) {
     std::lock_guard<std::mutex> lock(mMutex);
 
     // Check if state tracker already exists
-    if (mStateTrackers.find(stateAtomId) == mStateTrackers.end()) {
+    if (mStateTrackers.find(atomId) == mStateTrackers.end()) {
         // Create a new state tracker iff atom is a state atom
-        auto it = android::util::AtomsInfo::kStateAtomsFieldOptions.find(stateAtomId);
+        auto it = android::util::AtomsInfo::kStateAtomsFieldOptions.find(atomId);
         if (it != android::util::AtomsInfo::kStateAtomsFieldOptions.end()) {
-            mStateTrackers[stateAtomId] = new StateTracker(stateAtomId, it->second);
+            mStateTrackers[atomId] = new StateTracker(atomId, it->second);
         } else {
-            ALOGE("StateManager cannot register listener, Atom %d is not a state atom",
-                  stateAtomId);
+            ALOGE("StateManager cannot register listener, Atom %d is not a state atom", atomId);
             return false;
         }
     }
-    mStateTrackers[stateAtomId]->registerListener(listener);
+    mStateTrackers[atomId]->registerListener(listener);
     return true;
 }
 
-void StateManager::unregisterListener(int stateAtomId, wp<StateListener> listener) {
+void StateManager::unregisterListener(int atomId, wp<StateListener> listener) {
     std::unique_lock<std::mutex> lock(mMutex);
 
     // Hold the sp<> until the lock is released so that ~StateTracker() is
@@ -62,7 +61,7 @@
     sp<StateTracker> toRemove;
 
     // Unregister listener from correct StateTracker
-    auto it = mStateTrackers.find(stateAtomId);
+    auto it = mStateTrackers.find(atomId);
     if (it != mStateTrackers.end()) {
         it->second->unregisterListener(listener);
 
@@ -73,15 +72,15 @@
         }
     } else {
         ALOGE("StateManager cannot unregister listener, StateTracker for atom %d does not exist",
-              stateAtomId);
+              atomId);
     }
     lock.unlock();
 }
 
-int StateManager::getState(int stateAtomId, const HashableDimensionKey& key) {
+int StateManager::getStateValue(int atomId, const HashableDimensionKey& key) {
     std::lock_guard<std::mutex> lock(mMutex);
-    if (mStateTrackers.find(stateAtomId) != mStateTrackers.end()) {
-        return mStateTrackers[stateAtomId]->getState(key);
+    if (mStateTrackers.find(atomId) != mStateTrackers.end()) {
+        return mStateTrackers[atomId]->getStateValue(key);
     }
 
     return StateTracker::kStateUnknown;
diff --git a/cmds/statsd/src/state/StateManager.h b/cmds/statsd/src/state/StateManager.h
index ce60f14..89ee6c0 100644
--- a/cmds/statsd/src/state/StateManager.h
+++ b/cmds/statsd/src/state/StateManager.h
@@ -15,10 +15,11 @@
  */
 #pragma once
 
-//#include <utils/Log.h>
+#include <gtest/gtest_prod.h>
+#include <inttypes.h>
 #include <utils/RefBase.h>
-#include "HashableDimensionKey.h"
 
+#include "HashableDimensionKey.h"
 #include "state/StateListener.h"
 #include "state/StateTracker.h"
 
@@ -38,29 +39,29 @@
     // Notifies the correct StateTracker of an event.
     void onLogEvent(const LogEvent& event);
 
-    // Returns true if stateAtomId is the id of a state atom and notifies the
-    // correct StateTracker to register the listener. If the correct
-    // StateTracker does not exist, a new StateTracker is created.
-    bool registerListener(int stateAtomId, wp<StateListener> listener);
+    // Returns true if atomId is being tracked and is associated with a state
+    // atom. StateManager notifies the correct StateTracker to register listener.
+    // If the correct StateTracker does not exist, a new StateTracker is created.
+    bool registerListener(int atomId, wp<StateListener> listener);
 
     // Notifies the correct StateTracker to unregister a listener
     // and removes the tracker if it no longer has any listeners.
-    void unregisterListener(int stateAtomId, wp<StateListener> listener);
+    void unregisterListener(int atomId, wp<StateListener> listener);
 
-    // Queries the correct StateTracker for the state that is mapped to the given
-    // query key.
+    // Queries the correct StateTracker for the original/un-mapped state value
+    // that is mapped to the given query key.
     // If the StateTracker doesn't exist, returns StateTracker::kStateUnknown.
-    int getState(int stateAtomId, const HashableDimensionKey& queryKey);
+    int getStateValue(int atomId, const HashableDimensionKey& queryKey);
 
     inline int getStateTrackersCount() {
         std::lock_guard<std::mutex> lock(mMutex);
         return mStateTrackers.size();
     }
 
-    inline int getListenersCount(int stateAtomId) {
+    inline int getListenersCount(int atomId) {
         std::lock_guard<std::mutex> lock(mMutex);
-        if (mStateTrackers.find(stateAtomId) != mStateTrackers.end()) {
-            return mStateTrackers[stateAtomId]->getListenersCount();
+        if (mStateTrackers.find(atomId) != mStateTrackers.end()) {
+            return mStateTrackers[atomId]->getListenersCount();
         }
         return -1;
     }
diff --git a/cmds/statsd/src/state/StateTracker.cpp b/cmds/statsd/src/state/StateTracker.cpp
index 5a91950..323fc0e 100644
--- a/cmds/statsd/src/state/StateTracker.cpp
+++ b/cmds/statsd/src/state/StateTracker.cpp
@@ -30,13 +30,13 @@
   : mAtomId(atomId),
     mStateField(getSimpleMatcher(atomId, stateAtomInfo.exclusiveField)) {
     // create matcher for each primary field
-    // TODO(tsaichristine): handle when primary field is first uid in chain
+    // TODO(tsaichristine): b/142108433 handle when primary field is first uid in chain
     for (const auto& primary : stateAtomInfo.primaryFields) {
         Matcher matcher = getSimpleMatcher(atomId, primary);
         mPrimaryFields.push_back(matcher);
     }
 
-    // TODO(tsaichristine): set default state, reset state, and nesting
+    // TODO(tsaichristine): b/142108433 set default state, reset state, and nesting
 }
 
 void StateTracker::onLogEvent(const LogEvent& event) {
@@ -96,7 +96,7 @@
     mListeners.erase(listener);
 }
 
-int StateTracker::getState(const HashableDimensionKey& queryKey) const {
+int StateTracker::getStateValue(const HashableDimensionKey& queryKey) const {
     if (queryKey.getValues().size() == mPrimaryFields.size()) {
         auto it = mStateMap.find(queryKey);
         if (it != mStateMap.end()) {
diff --git a/cmds/statsd/src/state/StateTracker.h b/cmds/statsd/src/state/StateTracker.h
index f22706c..cfa9fd8 100644
--- a/cmds/statsd/src/state/StateTracker.h
+++ b/cmds/statsd/src/state/StateTracker.h
@@ -48,7 +48,7 @@
     // Returns the state value mapped to the given query key.
     // If the key isn't mapped to a state or the key size doesn't match the
     // primary key size, the default state is returned.
-    int getState(const HashableDimensionKey& queryKey) const;
+    int getStateValue(const HashableDimensionKey& queryKey) const;
 
     inline int getListenersCount() const {
         return mListeners.size();
diff --git a/cmds/statsd/tests/e2e/CountMetric_e2e_test.cpp b/cmds/statsd/tests/e2e/CountMetric_e2e_test.cpp
new file mode 100644
index 0000000..6591d69
--- /dev/null
+++ b/cmds/statsd/tests/e2e/CountMetric_e2e_test.cpp
@@ -0,0 +1,197 @@
+/*
+ * 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 <gtest/gtest.h>
+
+#include "src/StatsLogProcessor.h"
+#include "src/state/StateManager.h"
+#include "tests/statsd_test_util.h"
+
+namespace android {
+namespace os {
+namespace statsd {
+
+#ifdef __ANDROID__
+
+TEST(CountMetricE2eTest, TestWithSimpleState) {
+    // Initialize config
+    StatsdConfig config;
+    config.add_allowed_log_source("AID_ROOT");  // LogEvent defaults to UID of root.
+
+    auto syncStartMatcher = CreateSyncStartAtomMatcher();
+    *config.add_atom_matcher() = syncStartMatcher;
+
+    auto state = CreateScreenState();
+    *config.add_state() = state;
+
+    // Create count metric that slices by screen state
+    int64_t metricId = 123456;
+    auto countMetric = config.add_count_metric();
+    countMetric->set_id(metricId);
+    countMetric->set_what(syncStartMatcher.id());
+    countMetric->set_bucket(TimeUnit::ONE_MINUTE);
+    countMetric->add_slice_by_state(state.id());
+
+    // Initialize StatsLogProcessor
+    const int64_t baseTimeNs = 0;                                   // 0:00
+    const int64_t configAddedTimeNs = baseTimeNs + 1 * NS_PER_SEC;  // 0:01
+    const int64_t bucketSizeNs =
+            TimeUnitToBucketSizeInMillis(config.count_metric(0).bucket()) * 1000LL * 1000LL;
+
+    int uid = 12345;
+    int64_t cfgId = 98765;
+    ConfigKey cfgKey(uid, cfgId);
+
+    auto processor = CreateStatsLogProcessor(baseTimeNs, configAddedTimeNs, config, cfgKey);
+
+    // Check that StateTrackers were properly initialized.
+    EXPECT_EQ(1, StateManager::getInstance().getStateTrackersCount());
+    EXPECT_EQ(1,
+              StateManager::getInstance().getListenersCount(android::util::SCREEN_STATE_CHANGED));
+
+    // Check that CountMetricProducer was initialized correctly.
+    EXPECT_EQ(processor->mMetricsManagers.size(), 1u);
+    sp<MetricsManager> metricsManager = processor->mMetricsManagers.begin()->second;
+    EXPECT_TRUE(metricsManager->isConfigValid());
+    EXPECT_EQ(metricsManager->mAllMetricProducers.size(), 1);
+    sp<MetricProducer> metricProducer = metricsManager->mAllMetricProducers[0];
+    EXPECT_EQ(metricProducer->mSlicedStateAtoms.size(), 1);
+    EXPECT_EQ(metricProducer->mSlicedStateAtoms.at(0), android::util::SCREEN_STATE_CHANGED);
+    EXPECT_EQ(metricProducer->mStateGroupMap.size(), 0);
+}
+
+TEST(CountMetricE2eTest, TestWithMappedState) {
+    // Initialize config
+    StatsdConfig config;
+    config.add_allowed_log_source("AID_ROOT");  // LogEvent defaults to UID of root.
+
+    auto syncStartMatcher = CreateSyncStartAtomMatcher();
+    *config.add_atom_matcher() = syncStartMatcher;
+
+    auto state = CreateScreenStateWithOnOffMap();
+    *config.add_state() = state;
+
+    // Create count metric that slices by screen state with on/off map
+    int64_t metricId = 123456;
+    auto countMetric = config.add_count_metric();
+    countMetric->set_id(metricId);
+    countMetric->set_what(syncStartMatcher.id());
+    countMetric->set_bucket(TimeUnit::ONE_MINUTE);
+    countMetric->add_slice_by_state(state.id());
+
+    // Initialize StatsLogProcessor
+    const int64_t baseTimeNs = 0;                                   // 0:00
+    const int64_t configAddedTimeNs = baseTimeNs + 1 * NS_PER_SEC;  // 0:01
+    const int64_t bucketSizeNs =
+            TimeUnitToBucketSizeInMillis(config.count_metric(0).bucket()) * 1000LL * 1000LL;
+
+    int uid = 12345;
+    int64_t cfgId = 98765;
+    ConfigKey cfgKey(uid, cfgId);
+
+    auto processor = CreateStatsLogProcessor(baseTimeNs, configAddedTimeNs, config, cfgKey);
+
+    // Check that StateTrackers were properly initialized.
+    EXPECT_EQ(1, StateManager::getInstance().getStateTrackersCount());
+    EXPECT_EQ(1,
+              StateManager::getInstance().getListenersCount(android::util::SCREEN_STATE_CHANGED));
+
+    // Check that CountMetricProducer was initialized correctly.
+    EXPECT_EQ(processor->mMetricsManagers.size(), 1u);
+    sp<MetricsManager> metricsManager = processor->mMetricsManagers.begin()->second;
+    EXPECT_TRUE(metricsManager->isConfigValid());
+    EXPECT_EQ(metricsManager->mAllMetricProducers.size(), 1);
+    sp<MetricProducer> metricProducer = metricsManager->mAllMetricProducers[0];
+    EXPECT_EQ(metricProducer->mSlicedStateAtoms.size(), 1);
+    EXPECT_EQ(metricProducer->mSlicedStateAtoms.at(0), android::util::SCREEN_STATE_CHANGED);
+    EXPECT_EQ(metricProducer->mStateGroupMap.size(), 1);
+
+    StateMap map = state.map();
+    for (auto group : map.group()) {
+        for (auto value : group.value()) {
+            EXPECT_EQ(metricProducer->mStateGroupMap[android::util::SCREEN_STATE_CHANGED][value],
+                      group.group_id());
+        }
+    }
+}
+
+TEST(CountMetricE2eTest, TestWithMultipleStates) {
+    // Initialize config
+    StatsdConfig config;
+    config.add_allowed_log_source("AID_ROOT");  // LogEvent defaults to UID of root.
+
+    auto syncStartMatcher = CreateSyncStartAtomMatcher();
+    *config.add_atom_matcher() = syncStartMatcher;
+
+    auto state1 = CreateScreenStateWithOnOffMap();
+    *config.add_state() = state1;
+    auto state2 = CreateUidProcessState();
+    *config.add_state() = state2;
+
+    // Create count metric that slices by screen state with on/off map
+    int64_t metricId = 123456;
+    auto countMetric = config.add_count_metric();
+    countMetric->set_id(metricId);
+    countMetric->set_what(syncStartMatcher.id());
+    countMetric->set_bucket(TimeUnit::ONE_MINUTE);
+    countMetric->add_slice_by_state(state1.id());
+    countMetric->add_slice_by_state(state2.id());
+
+    // Initialize StatsLogProcessor
+    const int64_t baseTimeNs = 0;                                   // 0:00
+    const int64_t configAddedTimeNs = baseTimeNs + 1 * NS_PER_SEC;  // 0:01
+    const int64_t bucketSizeNs =
+            TimeUnitToBucketSizeInMillis(config.count_metric(0).bucket()) * 1000LL * 1000LL;
+
+    int uid = 12345;
+    int64_t cfgId = 98765;
+    ConfigKey cfgKey(uid, cfgId);
+
+    auto processor = CreateStatsLogProcessor(baseTimeNs, configAddedTimeNs, config, cfgKey);
+
+    // Check that StateTrackers were properly initialized.
+    EXPECT_EQ(2, StateManager::getInstance().getStateTrackersCount());
+    EXPECT_EQ(1,
+              StateManager::getInstance().getListenersCount(android::util::SCREEN_STATE_CHANGED));
+    EXPECT_EQ(1, StateManager::getInstance().getListenersCount(
+                         android::util::UID_PROCESS_STATE_CHANGED));
+
+    // Check that CountMetricProducer was initialized correctly.
+    EXPECT_EQ(processor->mMetricsManagers.size(), 1u);
+    sp<MetricsManager> metricsManager = processor->mMetricsManagers.begin()->second;
+    EXPECT_TRUE(metricsManager->isConfigValid());
+    EXPECT_EQ(metricsManager->mAllMetricProducers.size(), 1);
+    sp<MetricProducer> metricProducer = metricsManager->mAllMetricProducers[0];
+    EXPECT_EQ(metricProducer->mSlicedStateAtoms.size(), 2);
+    EXPECT_EQ(metricProducer->mSlicedStateAtoms.at(0), android::util::SCREEN_STATE_CHANGED);
+    EXPECT_EQ(metricProducer->mSlicedStateAtoms.at(1), android::util::UID_PROCESS_STATE_CHANGED);
+    EXPECT_EQ(metricProducer->mStateGroupMap.size(), 1);
+
+    StateMap map = state1.map();
+    for (auto group : map.group()) {
+        for (auto value : group.value()) {
+            EXPECT_EQ(metricProducer->mStateGroupMap[android::util::SCREEN_STATE_CHANGED][value],
+                      group.group_id());
+        }
+    }
+}
+
+}  // namespace statsd
+}  // namespace os
+}  // namespace android
+#else
+GTEST_LOG_(INFO) << "This test does nothing.\n";
+#endif
diff --git a/cmds/statsd/tests/metrics/CountMetricProducer_test.cpp b/cmds/statsd/tests/metrics/CountMetricProducer_test.cpp
index c503234..8915c73 100644
--- a/cmds/statsd/tests/metrics/CountMetricProducer_test.cpp
+++ b/cmds/statsd/tests/metrics/CountMetricProducer_test.cpp
@@ -43,8 +43,8 @@
     metric.set_bucket(ONE_MINUTE);
     sp<MockConditionWizard> wizard = new NaggyMock<MockConditionWizard>();
 
-    CountMetricProducer countProducer(kConfigKey, metric, -1 /*-1 meaning no condition*/, wizard,
-                                      5, 600 * NS_PER_SEC + NS_PER_SEC/2);
+    CountMetricProducer countProducer(kConfigKey, metric, -1 /*-1 meaning no condition*/, wizard, 5,
+                                      600 * NS_PER_SEC + NS_PER_SEC / 2);
     EXPECT_EQ(600500000000, countProducer.mCurrentBucketStartTimeNs);
     EXPECT_EQ(10, countProducer.mCurrentBucketNum);
     EXPECT_EQ(660000000005, countProducer.getCurrentBucketEndTimeNs());
@@ -131,7 +131,8 @@
 
     sp<MockConditionWizard> wizard = new NaggyMock<MockConditionWizard>();
 
-    CountMetricProducer countProducer(kConfigKey, metric, 1, wizard, bucketStartTimeNs, bucketStartTimeNs);
+    CountMetricProducer countProducer(kConfigKey, metric, 1, wizard, bucketStartTimeNs,
+                                      bucketStartTimeNs);
 
     countProducer.onConditionChanged(true, bucketStartTimeNs);
     countProducer.onMatchedLogEvent(1 /*matcher index*/, event1);
diff --git a/cmds/statsd/tests/state/StateTracker_test.cpp b/cmds/statsd/tests/state/StateTracker_test.cpp
index c89ffea..8d38000 100644
--- a/cmds/statsd/tests/state/StateTracker_test.cpp
+++ b/cmds/statsd/tests/state/StateTracker_test.cpp
@@ -44,7 +44,7 @@
 
     std::vector<Update> updates;
 
-    void onStateChanged(int stateAtomId, const HashableDimensionKey& primaryKey, int oldState,
+    void onStateChanged(int atomId, const HashableDimensionKey& primaryKey, int oldState,
                         int newState) {
         updates.emplace_back(primaryKey, newState);
     }
@@ -82,6 +82,7 @@
     return event;
 }
 
+// Incorrect event - missing fields
 std::shared_ptr<LogEvent> buildIncorrectOverlayEvent(int uid, const std::string& packageName, int state) {
     std::shared_ptr<LogEvent> event =
             std::make_shared<LogEvent>(android::util::OVERLAY_STATE_CHANGED, 1000 /*timestamp*/);
@@ -91,6 +92,18 @@
     event->init();
     return event;
 }
+
+// Incorrect event - exclusive state has wrong type
+std::shared_ptr<LogEvent> buildOverlayEventBadStateType(int uid, const std::string& packageName) {
+    std::shared_ptr<LogEvent> event =
+            std::make_shared<LogEvent>(android::util::OVERLAY_STATE_CHANGED, 1000 /*timestamp*/);
+    event->write((int32_t)uid);
+    event->write(packageName);
+    event->write(true);
+    event->write("string");  // exclusive state: string instead of int
+    event->init();
+    return event;
+}
 // END: build event functions.
 
 // START: get primary key functions
@@ -148,22 +161,22 @@
 
     // Register listener to non-existing StateTracker
     EXPECT_EQ(0, mgr.getStateTrackersCount());
-    mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener1);
+    EXPECT_TRUE(mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener1));
     EXPECT_EQ(1, mgr.getStateTrackersCount());
     EXPECT_EQ(1, mgr.getListenersCount(android::util::SCREEN_STATE_CHANGED));
 
     // Register listener to existing StateTracker
-    mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener2);
+    EXPECT_TRUE(mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener2));
     EXPECT_EQ(1, mgr.getStateTrackersCount());
     EXPECT_EQ(2, mgr.getListenersCount(android::util::SCREEN_STATE_CHANGED));
 
     // Register already registered listener to existing StateTracker
-    mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener2);
+    EXPECT_TRUE(mgr.registerListener(android::util::SCREEN_STATE_CHANGED, listener2));
     EXPECT_EQ(1, mgr.getStateTrackersCount());
     EXPECT_EQ(2, mgr.getListenersCount(android::util::SCREEN_STATE_CHANGED));
 
     // Register listener to non-state atom
-    mgr.registerListener(android::util::BATTERY_LEVEL_CHANGED, listener2);
+    EXPECT_FALSE(mgr.registerListener(android::util::BATTERY_LEVEL_CHANGED, listener2));
     EXPECT_EQ(1, mgr.getStateTrackersCount());
 }
 
@@ -227,12 +240,12 @@
 
     // check StateTracker was updated by querying for state
     HashableDimensionKey queryKey = DEFAULT_DIMENSION_KEY;
-    EXPECT_EQ(2, mgr.getState(android::util::SCREEN_STATE_CHANGED, queryKey));
+    EXPECT_EQ(2, mgr.getStateValue(android::util::SCREEN_STATE_CHANGED, queryKey));
 }
 
 /**
  * Test StateManager's onLogEvent and StateListener's onStateChanged correctly
- * updates listener for states with primary keys.
+ * updates listener for states with one primary key.
  */
 TEST(StateTrackerTest, TestStateChangeOnePrimaryField) {
     sp<TestStateListener> listener1 = new TestStateListener();
@@ -240,9 +253,8 @@
     mgr.registerListener(android::util::UID_PROCESS_STATE_CHANGED, listener1);
 
     // log event
-    std::shared_ptr<LogEvent> event = buildUidProcessEvent(
-            1000,
-            android::app::ProcessStateEnum::PROCESS_STATE_TOP);  //  state value: 1002
+    std::shared_ptr<LogEvent> event =
+            buildUidProcessEvent(1000 /* uid */, android::app::ProcessStateEnum::PROCESS_STATE_TOP);
     mgr.onLogEvent(*event);
 
     // check listener was updated
@@ -252,23 +264,33 @@
 
     // check StateTracker was updated by querying for state
     HashableDimensionKey queryKey;
-    getUidProcessKey(1000, &queryKey);
-    EXPECT_EQ(1002, mgr.getState(android::util::UID_PROCESS_STATE_CHANGED, queryKey));
+    getUidProcessKey(1000 /* uid */, &queryKey);
+    EXPECT_EQ(1002, mgr.getStateValue(android::util::UID_PROCESS_STATE_CHANGED, queryKey));
 }
 
+/**
+ * Test StateManager's onLogEvent and StateListener's onStateChanged correctly
+ * updates listener for states with multiple primary keys.
+ */
 TEST(StateTrackerTest, TestStateChangeMultiplePrimaryFields) {
     sp<TestStateListener> listener1 = new TestStateListener();
     StateManager mgr;
     mgr.registerListener(android::util::OVERLAY_STATE_CHANGED, listener1);
 
     // log event
-    std::shared_ptr<LogEvent> event = buildOverlayEvent(1000, "package1", 1);  // state: ENTERED
+    std::shared_ptr<LogEvent> event =
+            buildOverlayEvent(1000 /* uid */, "package1", 1);  // state: ENTERED
     mgr.onLogEvent(*event);
 
-    // check listener update
+    // check listener was updated
     EXPECT_EQ(1, listener1->updates.size());
     EXPECT_EQ(1000, listener1->updates[0].mKey.getValues()[0].mValue.int_value);
     EXPECT_EQ(1, listener1->updates[0].mState);
+
+    // check StateTracker was updated by querying for state
+    HashableDimensionKey queryKey;
+    getOverlayKey(1000 /* uid */, "package1", &queryKey);
+    EXPECT_EQ(1, mgr.getStateValue(android::util::OVERLAY_STATE_CHANGED, queryKey));
 }
 
 /**
@@ -282,11 +304,14 @@
     mgr.registerListener(android::util::OVERLAY_STATE_CHANGED, listener1);
 
     // log event
-    std::shared_ptr<LogEvent> event =
-            buildIncorrectOverlayEvent(1000, "package1", 1);  // state: ENTERED
-    mgr.onLogEvent(*event);
+    std::shared_ptr<LogEvent> event1 =
+            buildIncorrectOverlayEvent(1000 /* uid */, "package1", 1 /* state */);
+    std::shared_ptr<LogEvent> event2 = buildOverlayEventBadStateType(1001 /* uid */, "package2");
 
-    // check listener update
+    // check listener was updated
+    mgr.onLogEvent(*event1);
+    EXPECT_EQ(0, listener1->updates.size());
+    mgr.onLogEvent(*event2);
     EXPECT_EQ(0, listener1->updates.size());
 }
 
@@ -301,18 +326,19 @@
 
     std::shared_ptr<LogEvent> event1 = buildUidProcessEvent(
             1000,
-            android::app::ProcessStateEnum::PROCESS_STATE_TOP);  // state value: 1002
+            android::app::ProcessStateEnum::PROCESS_STATE_TOP);  //  state value: 1002
     std::shared_ptr<LogEvent> event2 = buildUidProcessEvent(
             1001,
-            android::app::ProcessStateEnum::PROCESS_STATE_FOREGROUND_SERVICE);  // state value: 1003
+            android::app::ProcessStateEnum::PROCESS_STATE_FOREGROUND_SERVICE);  //  state value:
+                                                                                //  1003
     std::shared_ptr<LogEvent> event3 = buildUidProcessEvent(
             1002,
-            android::app::ProcessStateEnum::PROCESS_STATE_PERSISTENT);  // state value: 1000
+            android::app::ProcessStateEnum::PROCESS_STATE_PERSISTENT);  //  state value: 1000
     std::shared_ptr<LogEvent> event4 = buildUidProcessEvent(
             1001,
-            android::app::ProcessStateEnum::PROCESS_STATE_TOP);  // state value: 1002
+            android::app::ProcessStateEnum::PROCESS_STATE_TOP);  //  state value: 1002
     std::shared_ptr<LogEvent> event5 =
-            buildScreenEvent(android::view::DisplayStateEnum::DISPLAY_STATE_ON);  // state value:
+            buildScreenEvent(android::view::DisplayStateEnum::DISPLAY_STATE_ON);
     std::shared_ptr<LogEvent> event6 = buildOverlayEvent(1000, "package1", 1);
     std::shared_ptr<LogEvent> event7 = buildOverlayEvent(1000, "package2", 2);
 
@@ -327,25 +353,25 @@
     // Query for UidProcessState of uid 1001
     HashableDimensionKey queryKey1;
     getUidProcessKey(1001, &queryKey1);
-    EXPECT_EQ(1003, mgr.getState(android::util::UID_PROCESS_STATE_CHANGED, queryKey1));
+    EXPECT_EQ(1003, mgr.getStateValue(android::util::UID_PROCESS_STATE_CHANGED, queryKey1));
 
     // Query for UidProcessState of uid 1004 - not in state map
     HashableDimensionKey queryKey2;
     getUidProcessKey(1004, &queryKey2);
-    EXPECT_EQ(-1,
-              mgr.getState(android::util::UID_PROCESS_STATE_CHANGED, queryKey2));  // default state
+    EXPECT_EQ(-1, mgr.getStateValue(android::util::UID_PROCESS_STATE_CHANGED,
+                                    queryKey2));  // default state
 
     // Query for UidProcessState of uid 1001 - after change in state
     mgr.onLogEvent(*event4);
-    EXPECT_EQ(1002, mgr.getState(android::util::UID_PROCESS_STATE_CHANGED, queryKey1));
+    EXPECT_EQ(1002, mgr.getStateValue(android::util::UID_PROCESS_STATE_CHANGED, queryKey1));
 
     // Query for ScreenState
-    EXPECT_EQ(2, mgr.getState(android::util::SCREEN_STATE_CHANGED, DEFAULT_DIMENSION_KEY));
+    EXPECT_EQ(2, mgr.getStateValue(android::util::SCREEN_STATE_CHANGED, DEFAULT_DIMENSION_KEY));
 
     // Query for OverlayState of uid 1000, package name "package2"
     HashableDimensionKey queryKey3;
     getOverlayKey(1000, "package2", &queryKey3);
-    EXPECT_EQ(2, mgr.getState(android::util::OVERLAY_STATE_CHANGED, queryKey3));
+    EXPECT_EQ(2, mgr.getStateValue(android::util::OVERLAY_STATE_CHANGED, queryKey3));
 }
 
 }  // namespace statsd
diff --git a/cmds/statsd/tests/statsd_test_util.cpp b/cmds/statsd/tests/statsd_test_util.cpp
index 2c4f3c7..38c22ab 100644
--- a/cmds/statsd/tests/statsd_test_util.cpp
+++ b/cmds/statsd/tests/statsd_test_util.cpp
@@ -251,6 +251,101 @@
     return predicate;
 }
 
+State CreateScreenState() {
+    State state;
+    state.set_id(StringToId("ScreenState"));
+    state.set_atom_id(29);
+    return state;
+}
+
+State CreateUidProcessState() {
+    State state;
+    state.set_id(StringToId("UidProcessState"));
+    state.set_atom_id(27);
+    return state;
+}
+
+State CreateOverlayState() {
+    State state;
+    state.set_id(StringToId("OverlayState"));
+    state.set_atom_id(59);
+    return state;
+}
+
+State CreateScreenStateWithOnOffMap() {
+    State state;
+    state.set_id(StringToId("ScreenStateOnOff"));
+    state.set_atom_id(29);
+
+    auto map = CreateScreenStateOnOffMap();
+    *state.mutable_map() = map;
+
+    return state;
+}
+
+State CreateScreenStateWithInDozeMap() {
+    State state;
+    state.set_id(StringToId("ScreenStateInDoze"));
+    state.set_atom_id(29);
+
+    auto map = CreateScreenStateInDozeMap();
+    *state.mutable_map() = map;
+
+    return state;
+}
+
+StateMap_StateGroup CreateScreenStateOnGroup() {
+    StateMap_StateGroup group;
+    group.set_group_id(StringToId("SCREEN_ON"));
+    group.add_value(2);
+    group.add_value(5);
+    group.add_value(6);
+    return group;
+}
+
+StateMap_StateGroup CreateScreenStateOffGroup() {
+    StateMap_StateGroup group;
+    group.set_group_id(StringToId("SCREEN_OFF"));
+    group.add_value(0);
+    group.add_value(1);
+    group.add_value(3);
+    group.add_value(4);
+    return group;
+}
+
+StateMap CreateScreenStateOnOffMap() {
+    StateMap map;
+    *map.add_group() = CreateScreenStateOnGroup();
+    *map.add_group() = CreateScreenStateOffGroup();
+    return map;
+}
+
+StateMap_StateGroup CreateScreenStateInDozeGroup() {
+    StateMap_StateGroup group;
+    group.set_group_id(StringToId("SCREEN_DOZE"));
+    group.add_value(3);
+    group.add_value(4);
+    return group;
+}
+
+StateMap_StateGroup CreateScreenStateNotDozeGroup() {
+    StateMap_StateGroup group;
+    group.set_group_id(StringToId("SCREEN_NOT_DOZE"));
+    group.add_value(0);
+    group.add_value(1);
+    group.add_value(2);
+    group.add_value(5);
+    group.add_value(6);
+    return group;
+}
+
+StateMap CreateScreenStateInDozeMap() {
+    StateMap map;
+    *map.add_group() = CreateScreenStateInDozeGroup();
+    *map.add_group() = CreateScreenStateNotDozeGroup();
+    return map;
+}
+
 void addPredicateToPredicateCombination(const Predicate& predicate,
                                         Predicate* combinationPredicate) {
     combinationPredicate->mutable_combination()->add_predicate(predicate.id());
diff --git a/cmds/statsd/tests/statsd_test_util.h b/cmds/statsd/tests/statsd_test_util.h
index 635c583..c026105 100644
--- a/cmds/statsd/tests/statsd_test_util.h
+++ b/cmds/statsd/tests/statsd_test_util.h
@@ -104,6 +104,37 @@
 // Create a Predicate proto for app is in background.
 Predicate CreateIsInBackgroundPredicate();
 
+// Create State proto for screen state atom.
+State CreateScreenState();
+
+// Create State proto for uid process state atom.
+State CreateUidProcessState();
+
+// Create State proto for overlay state atom.
+State CreateOverlayState();
+
+State CreateScreenStateWithOnOffMap();
+
+State CreateScreenStateWithInDozeMap();
+
+// Create StateGroup proto for ScreenState ON group
+StateMap_StateGroup CreateScreenStateOnGroup();
+
+// Create StateGroup proto for ScreenState OFF group
+StateMap_StateGroup CreateScreenStateOffGroup();
+
+// Create StateMap proto for ScreenState ON/OFF map
+StateMap CreateScreenStateOnOffMap();
+
+// Create StateGroup proto for ScreenState IN DOZE group
+StateMap_StateGroup CreateScreenStateInDozeGroup();
+
+// Create StateGroup proto for ScreenState NOT IN DOZE group
+StateMap_StateGroup CreateScreenStateNotDozeGroup();
+
+// Create StateMap proto for ScreenState IN DOZE map
+StateMap CreateScreenStateInDozeMap();
+
 // Add a predicate to the predicate combination.
 void addPredicateToPredicateCombination(const Predicate& predicate, Predicate* combination);
 
@@ -319,4 +350,4 @@
 }
 }  // namespace statsd
 }  // namespace os
-}  // namespace android
\ No newline at end of file
+}  // namespace android