diff --git a/server/TrafficControllerTest.cpp b/server/TrafficControllerTest.cpp
index 1ad3915..2d47985 100644
--- a/server/TrafficControllerTest.cpp
+++ b/server/TrafficControllerTest.cpp
@@ -33,6 +33,7 @@
 #include <android-base/strings.h>
 
 #include <netdutils/MockSyscalls.h>
+#include "FirewallController.h"
 #include "TrafficController.h"
 #include "bpf/BpfUtils.h"
 
@@ -53,6 +54,8 @@
 
 constexpr int TEST_MAP_SIZE = 10;
 constexpr uid_t TEST_UID = 10086;
+constexpr uid_t TEST_UID2 = 54321;
+constexpr uid_t TEST_UID3 = 98765;
 constexpr uint32_t TEST_TAG = 42;
 constexpr int TEST_COUNTERSET = 1;
 constexpr int DEFAULT_COUNTERSET = 0;
@@ -70,8 +73,12 @@
     unique_fd mFakeUidCounterSetMap;
     unique_fd mFakeUidStatsMap;
     unique_fd mFakeTagStatsMap;
+    unique_fd mFakeDozableUidMap;
+    unique_fd mFakeStandbyUidMap;
+    unique_fd mFakePowerSaveUidMap;
 
     void SetUp() {
+        std::lock_guard<std::mutex> guard(mTc.mOwnerMatchMutex);
         SKIP_IF_BPF_NOT_SUPPORTED;
 
         mFakeCookieTagMap = unique_fd(createMap(BPF_MAP_TYPE_HASH, sizeof(uint64_t),
@@ -90,6 +97,17 @@
                                                sizeof(struct StatsValue), TEST_MAP_SIZE, 0));
         ASSERT_LE(0, mFakeTagStatsMap);
 
+        mFakeDozableUidMap = unique_fd(
+            createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint32_t), TEST_MAP_SIZE, 0));
+        ASSERT_LE(0, mFakeDozableUidMap);
+
+        mFakeStandbyUidMap = unique_fd(
+            createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint32_t), TEST_MAP_SIZE, 0));
+        ASSERT_LE(0, mFakeStandbyUidMap);
+
+        mFakePowerSaveUidMap = unique_fd(
+            createMap(BPF_MAP_TYPE_HASH, sizeof(uint32_t), sizeof(uint32_t), TEST_MAP_SIZE, 0));
+        ASSERT_LE(0, mFakePowerSaveUidMap);
         // Make sure trafficController use the eBPF code path.
         mTc.ebpfSupported = true;
 
@@ -97,6 +115,9 @@
         mTc.mUidCounterSetMap.reset(mFakeUidCounterSetMap);
         mTc.mUidStatsMap.reset(mFakeUidStatsMap);
         mTc.mTagStatsMap.reset(mFakeTagStatsMap);
+        mTc.mDozableUidMap.reset(mFakeDozableUidMap);
+        mTc.mStandbyUidMap.reset(mFakeStandbyUidMap);
+        mTc.mPowerSaveUidMap.reset(mFakePowerSaveUidMap);
     }
 
     int setUpSocketAndTag(int protocol, uint64_t* cookie, uint32_t tag, uid_t uid) {
@@ -140,6 +161,62 @@
         key->tag = tag;
     }
 
+    void checkUidOwnerRuleForChain(ChildChain chain, const unique_fd& targetMap) {
+        uint32_t uid = TEST_UID;
+        EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, DENY, BLACKLIST));
+        uint8_t value;
+        EXPECT_EQ(0, findMapEntry(targetMap, &uid, &value));
+        EXPECT_EQ((uint8_t)BPF_DROP, value);
+
+        uid = TEST_UID2;
+        EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, ALLOW, WHITELIST));
+        EXPECT_EQ(0, findMapEntry(targetMap, &uid, &value));
+        EXPECT_EQ((uint8_t)BPF_PASS, value);
+
+        EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, DENY, WHITELIST));
+        EXPECT_EQ(-1, findMapEntry(targetMap, &uid, &value));
+        EXPECT_EQ(ENOENT, errno);
+
+        uid = TEST_UID;
+        EXPECT_EQ(0, mTc.changeUidOwnerRule(chain, uid, ALLOW, BLACKLIST));
+        EXPECT_EQ(-1, findMapEntry(targetMap, &uid, &value));
+        EXPECT_EQ(ENOENT, errno);
+
+        uid = TEST_UID3;
+        EXPECT_EQ(-ENOENT, mTc.changeUidOwnerRule(chain, uid, ALLOW, BLACKLIST));
+        EXPECT_EQ(-1, findMapEntry(targetMap, &uid, &value));
+        EXPECT_EQ(ENOENT, errno);
+    }
+
+    void checkEachUidValue(const std::vector<int32_t>& uids, const uint8_t expectValue,
+                           const unique_fd& targetMap) {
+        uint8_t value;
+        for (auto uid : uids) {
+            EXPECT_EQ(0, findMapEntry(targetMap, &uid, &value));
+            EXPECT_EQ((uint8_t)expectValue, value);
+        }
+        std::set<uint32_t> uidSet(uids.begin(), uids.end());
+        auto checkNoOtherUid = [&uidSet](void *key, const base::unique_fd&) {
+            int32_t uid = *(int32_t *)key;
+            EXPECT_NE(uidSet.end(), uidSet.find(uid));
+            return 0;
+        };
+        uint32_t nonExistentKey = NONEXISTENT_UID;
+        uint8_t dummyValue;
+        EXPECT_EQ(0, bpfIterateMap(nonExistentKey, dummyValue, targetMap, checkNoOtherUid));
+    }
+
+    void checkUidMapReplace(const std::string& name, const std::vector<int32_t>& uids,
+                            const unique_fd& targetMap) {
+        bool isWhitelist = true;
+        EXPECT_EQ(0, mTc.replaceUidOwnerMap(name, isWhitelist, uids));
+        checkEachUidValue(uids, BPF_PASS, targetMap);
+
+        isWhitelist = false;
+        EXPECT_EQ(0, mTc.replaceUidOwnerMap(name, isWhitelist, uids));
+        checkEachUidValue(uids, BPF_DROP, targetMap);
+    }
+
     void TearDown() {
         mFakeCookieTagMap.reset();
         mFakeUidCounterSetMap.reset();
@@ -357,5 +434,54 @@
     ASSERT_EQ((uint64_t)200, statsMapResult.rxBytes);
 }
 
+TEST_F(TrafficControllerTest, TestUpdateOwnerMapEntry) {
+    SKIP_IF_BPF_NOT_SUPPORTED;
+
+    uint32_t uid = TEST_UID;
+    ASSERT_EQ(0, mTc.updateOwnerMapEntry(mFakeDozableUidMap, uid, DENY, BLACKLIST));
+    uint8_t value;
+    ASSERT_EQ(0, findMapEntry(mFakeDozableUidMap, &uid, &value));
+    ASSERT_EQ((uint8_t)BPF_DROP, value);
+
+    uid = TEST_UID2;
+    ASSERT_EQ(0, mTc.updateOwnerMapEntry(mFakeDozableUidMap, uid, ALLOW, WHITELIST));
+    ASSERT_EQ(0, findMapEntry(mFakeDozableUidMap, &uid, &value));
+    ASSERT_EQ((uint8_t)BPF_PASS, value);
+
+    ASSERT_EQ(0, mTc.updateOwnerMapEntry(mFakeDozableUidMap, uid, DENY, WHITELIST));
+    ASSERT_EQ(-1, findMapEntry(mFakeDozableUidMap, &uid, &value));
+    ASSERT_EQ(ENOENT, errno);
+
+    uid = TEST_UID;
+    ASSERT_EQ(0, mTc.updateOwnerMapEntry(mFakeDozableUidMap, uid, ALLOW, BLACKLIST));
+    ASSERT_EQ(-1, findMapEntry(mFakeDozableUidMap, &uid, &value));
+    ASSERT_EQ(ENOENT, errno);
+
+    uid = TEST_UID3;
+    ASSERT_EQ(-ENOENT, mTc.updateOwnerMapEntry(mFakeDozableUidMap, uid, ALLOW, BLACKLIST));
+    ASSERT_EQ(-1, findMapEntry(mFakeDozableUidMap, &uid, &value));
+    ASSERT_EQ(ENOENT, errno);
+}
+
+TEST_F(TrafficControllerTest, TestChangeUidOwnerRule) {
+    SKIP_IF_BPF_NOT_SUPPORTED;
+
+    checkUidOwnerRuleForChain(DOZABLE, mFakeDozableUidMap);
+    checkUidOwnerRuleForChain(STANDBY, mFakeStandbyUidMap);
+    checkUidOwnerRuleForChain(POWERSAVE, mFakePowerSaveUidMap);
+    ASSERT_EQ(-EINVAL, mTc.changeUidOwnerRule(NONE, TEST_UID, ALLOW, WHITELIST));
+    ASSERT_EQ(-EINVAL, mTc.changeUidOwnerRule(INVALID_CHAIN, TEST_UID, ALLOW, WHITELIST));
+}
+
+TEST_F(TrafficControllerTest, TestReplaceUidOwnerMap) {
+    SKIP_IF_BPF_NOT_SUPPORTED;
+
+    std::vector<int32_t> uids = {TEST_UID, TEST_UID2, TEST_UID3};
+    checkUidMapReplace("fw_dozable", uids, mFakeDozableUidMap);
+    checkUidMapReplace("fw_standby", uids, mFakeStandbyUidMap);
+    checkUidMapReplace("fw_powersave", uids, mFakePowerSaveUidMap);
+    ASSERT_EQ(-EINVAL, mTc.replaceUidOwnerMap("unknow", true, uids));
+}
+
 }  // namespace net
 }  // namespace android