diff --git a/test/unittest/cpp_test/src/input_method_private_member_test.cpp b/test/unittest/cpp_test/src/input_method_private_member_test.cpp
index 4a6ed8f6c91b6f6a2486fbb3e00aea574d46d056..81de5ff33b6d5c2e3ec29ba37ac51db266327bc5 100644
--- a/test/unittest/cpp_test/src/input_method_private_member_test.cpp
+++ b/test/unittest/cpp_test/src/input_method_private_member_test.cpp
@@ -639,19 +639,26 @@ HWTEST_F(InputMethodPrivateMemberTest, III_TestRestoreInputMethod_001, TestSize.
     auto ret = service_->RestoreInputmethod(bundleName);
     EXPECT_EQ(ret, ErrorCode::ERROR_ENABLE_IME);
 
-    auto currentProp = InputMethodController::GetInstance()->GetCurrentInputMethod();
+    auto userId = service_->GetCallingUserId();
+    auto currentProp = ImeInfoInquirer::GetInstance().GetCurrentInputMethod(userId);
     ret = service_->RestoreInputmethod(currentProp->name);
     EXPECT_EQ(ret, ErrorCode::NO_ERROR);
 
     auto defaultIme = ImeInfoInquirer::GetInstance().GetDefaultIme();
     ret = service_->RestoreInputmethod(defaultIme.bundleName);
-    EXPECT_EQ(ret, ErrorCode::NO_ERROR);
+    EXPECT_TRUE(ret == ErrorCode::NO_ERROR ||ret == ErrorCode::ERROR_IMSA_REBOOT_OLD_IME_NOT_STOP);
 
     bundleName = "com.example.newTestIme";
-    ret = service_->RestoreInputmethod(bundleName);
-    EXPECT_EQ(ret, ErrorCode::ERROR_IMSA_REBOOT_OLD_IME_NOT_STOP);
+    EnabledStatus status = EnabledStatus::DISABLED;
+    ImeInfoInquirer::GetInstance().systemConfig_.enableInputMethodFeature = true;
+    service_->EnableIme(userId, bundleName, "", status);
+    ImeEnabledInfoManager::GetInstance().GetEnabledStateInner(userId, bundleName, status);
+    EXPECT_EQ(status, EnabledStatus::DISABLED);
+    service_->RestoreInputmethod(bundleName);
+    ImeEnabledInfoManager::GetInstance().GetEnabledStateInner(userId, bundleName, status);
+    EXPECT_EQ(status, EnabledStatus::BASIC_MODE);
+    ImeInfoInquirer::GetInstance().systemConfig_.enableInputMethodFeature = false;
 
-    auto userId = service_->GetCallingUserId();
     UserSessionManager::GetInstance().RemoveUserSession(userId);
     ret = service_->RestoreInputmethod(defaultIme.bundleName);
     EXPECT_EQ(ret, ErrorCode::ERROR_NULL_POINTER);