diff --git a/services/inputflinger/tests/InputMapperTest.cpp b/services/inputflinger/tests/InputMapperTest.cpp
index 5722444..e773f58 100644
--- a/services/inputflinger/tests/InputMapperTest.cpp
+++ b/services/inputflinger/tests/InputMapperTest.cpp
@@ -89,6 +89,13 @@
     }
 }
 
+void InputMapperUnitTest::setSwitchState(int32_t state, std::set<int32_t> switchCodes) {
+    for (const auto& switchCode : switchCodes) {
+        EXPECT_CALL(mMockEventHub, getSwitchState(EVENTHUB_ID, switchCode))
+                .WillRepeatedly(testing::Return(static_cast<int>(state)));
+    }
+}
+
 std::list<NotifyArgs> InputMapperUnitTest::process(int32_t type, int32_t code, int32_t value) {
     nsecs_t when = systemTime(SYSTEM_TIME_MONOTONIC);
     return process(when, type, code, value);
diff --git a/services/inputflinger/tests/InputMapperTest.h b/services/inputflinger/tests/InputMapperTest.h
index 5bd8cda..88057dc 100644
--- a/services/inputflinger/tests/InputMapperTest.h
+++ b/services/inputflinger/tests/InputMapperTest.h
@@ -58,6 +58,8 @@
 
     void setKeyCodeState(KeyState state, std::set<int> keyCodes);
 
+    void setSwitchState(int32_t state, std::set<int32_t> switchCodes);
+
     std::list<NotifyArgs> process(int32_t type, int32_t code, int32_t value);
     std::list<NotifyArgs> process(nsecs_t when, int32_t type, int32_t code, int32_t value);
 
diff --git a/services/inputflinger/tests/SwitchInputMapper_test.cpp b/services/inputflinger/tests/SwitchInputMapper_test.cpp
index 7b67c01..4020e78 100644
--- a/services/inputflinger/tests/SwitchInputMapper_test.cpp
+++ b/services/inputflinger/tests/SwitchInputMapper_test.cpp
@@ -25,39 +25,41 @@
 #include <linux/input-event-codes.h>
 
 #include "InputMapperTest.h"
+#include "TestConstants.h"
 
 namespace android {
 
-class SwitchInputMapperTest : public InputMapperTest {
+class SwitchInputMapperTest : public InputMapperUnitTest {
 protected:
+    void SetUp() override {
+        InputMapperUnitTest::SetUp();
+        createDevice();
+        mMapper = createInputMapper<SwitchInputMapper>(*mDeviceContext,
+                                                       mFakePolicy->getReaderConfiguration());
+    }
 };
 
 TEST_F(SwitchInputMapperTest, GetSources) {
-    SwitchInputMapper& mapper = constructAndAddMapper<SwitchInputMapper>();
-
-    ASSERT_EQ(uint32_t(AINPUT_SOURCE_SWITCH), mapper.getSources());
+    ASSERT_EQ(uint32_t(AINPUT_SOURCE_SWITCH), mMapper->getSources());
 }
 
 TEST_F(SwitchInputMapperTest, GetSwitchState) {
-    SwitchInputMapper& mapper = constructAndAddMapper<SwitchInputMapper>();
+    setSwitchState(1, {SW_LID});
+    ASSERT_EQ(1, mMapper->getSwitchState(AINPUT_SOURCE_ANY, SW_LID));
 
-    mFakeEventHub->setSwitchState(EVENTHUB_ID, SW_LID, 1);
-    ASSERT_EQ(1, mapper.getSwitchState(AINPUT_SOURCE_ANY, SW_LID));
-
-    mFakeEventHub->setSwitchState(EVENTHUB_ID, SW_LID, 0);
-    ASSERT_EQ(0, mapper.getSwitchState(AINPUT_SOURCE_ANY, SW_LID));
+    setSwitchState(0, {SW_LID});
+    ASSERT_EQ(0, mMapper->getSwitchState(AINPUT_SOURCE_ANY, SW_LID));
 }
 
 TEST_F(SwitchInputMapperTest, Process) {
-    SwitchInputMapper& mapper = constructAndAddMapper<SwitchInputMapper>();
     std::list<NotifyArgs> out;
-    out = process(mapper, ARBITRARY_TIME, READ_TIME, EV_SW, SW_LID, 1);
+    out = process(ARBITRARY_TIME, EV_SW, SW_LID, 1);
     ASSERT_TRUE(out.empty());
-    out = process(mapper, ARBITRARY_TIME, READ_TIME, EV_SW, SW_JACK_PHYSICAL_INSERT, 1);
+    out = process(ARBITRARY_TIME, EV_SW, SW_JACK_PHYSICAL_INSERT, 1);
     ASSERT_TRUE(out.empty());
-    out = process(mapper, ARBITRARY_TIME, READ_TIME, EV_SW, SW_HEADPHONE_INSERT, 0);
+    out = process(ARBITRARY_TIME, EV_SW, SW_HEADPHONE_INSERT, 0);
     ASSERT_TRUE(out.empty());
-    out = process(mapper, ARBITRARY_TIME, READ_TIME, EV_SYN, SYN_REPORT, 0);
+    out = process(ARBITRARY_TIME, EV_SYN, SYN_REPORT, 0);
 
     ASSERT_EQ(1u, out.size());
     const NotifySwitchArgs& args = std::get<NotifySwitchArgs>(*out.begin());