diff --git a/test/test_action.py b/test/test_action.py
index 57679cf..974a89d 100644
--- a/test/test_action.py
+++ b/test/test_action.py
@@ -13,6 +13,7 @@ class EventMock():
 @dataclass
 class ResponseMock():
     status: str
+    keys: Dict[str, str]
 
 
 @dataclass
@@ -134,7 +135,7 @@ def test__attach_event_filter(self) -> None:
     def test__send_action(self) -> None:
         # Test general AMI error
         self.__client_mock.send_action_result = FutureResponseMock(
-            ResponseMock("error"))
+            ResponseMock("error", {"Message": "Some error"}))
         result = self.__ae._ActionExecuter__send_action()
         self.assertFalse(result, "Expected result to be false")
 
@@ -145,7 +146,7 @@ def test__send_action(self) -> None:
 
         # Test successful action
         self.__client_mock.send_action_result = FutureResponseMock(
-            ResponseMock("Success"))
+            ResponseMock("Success", {"Message": "Success"}))
         result = self.__ae._ActionExecuter__send_action()
         self.assertTrue(result, "Expected result to be true")
 
diff --git a/test/test_client_wrapper.py b/test/test_client_wrapper.py
index 2b8565c..b5b891f 100644
--- a/test/test_client_wrapper.py
+++ b/test/test_client_wrapper.py
@@ -127,9 +127,8 @@ def test_login(self):
 
         self.__ami_client.login_response = FutureResponseMock(
             ResponseMock("Error", False))
-        self.assertRaisesRegex(
-            Exception,
-            "Unable to login AMI client",
+        self.assertRaises(
+            SystemExit,
             self.__client.login,
             "",
             "",