Refactor test_login_attemps_correctly_reported in test_powershell.py to address the changes in the flow of powershell and powershell client

monkey/tests/unit_tests/infection_monkey/exploit/test_powershell.py

@@ -157,65 +157,28 @@ def test_failed_monkey_execution(monkeypatch, powershell_exploiter):
 
 def test_login_attemps_correctly_reported(monkeypatch, powershell_exploiter):
     mock_client = MagicMock()
-    mock_client.get_host_architecture = lambda: WIN_ARCH_32
-    mock_client.copy_file = MagicMock(return_value=True)
-    mock_client.execute_cmd_as_detached_process = MagicMock(side_effect=Exception)
+    mock_client.return_value.get_host_architecture = lambda: WIN_ARCH_32
+    mock_client.return_value.copy_file = MagicMock(return_value=True)
 
-    def allow_ntlm(_, credentials: Credentials, auth_options: AuthOptions):
-        if credentials.username == USER_LIST[1] and credentials.secret == NT_HASH_LIST[1]:
-            return mock_client
+    # execute_cmd method will throw exceptions for 5 first calls.
+    # 6-th call doesn't throw an exception == credentials successful
+    execute_cmd_returns = [Exception, Exception, Exception, Exception, Exception, True]
+    mock_client.return_value.execute_cmd = MagicMock(side_effect=execute_cmd_returns)
 
-        raise TestAuthenticationError
-
-    mock_powershell_client = MagicMock(side_effect=allow_ntlm)
-    monkeypatch.setattr(powershell, "PowerShellClient", mock_powershell_client)
+    monkeypatch.setattr(powershell, "PowerShellClient", mock_client)
 
     powershell_exploiter.exploit_host()
 
-    assert {
-        "result": False,
-        "user": USER_LIST[1],
-        "password": None,
-        "lm_hash": "",
-        "ntlm_hash": "",
-        "ssh_key": "",
-    } in powershell_exploiter.exploit_attempts
-
-    assert {
-        "result": False,
-        "user": USER_LIST[1],
-        "password": PASSWORD_LIST[0],
-        "lm_hash": "",
-        "ntlm_hash": "",
-        "ssh_key": "",
-    } in powershell_exploiter.exploit_attempts
-
-    assert {
-        "result": False,
-        "user": USER_LIST[0],
-        "password": "",
-        "lm_hash": LM_HASH_LIST[0],
-        "ntlm_hash": "",
-        "ssh_key": "",
-    } in powershell_exploiter.exploit_attempts
-
-    assert {
-        "result": False,
-        "user": USER_LIST[1],
-        "password": "",
-        "lm_hash": "",
-        "ntlm_hash": NT_HASH_LIST[0],
-        "ssh_key": "",
-    } in powershell_exploiter.exploit_attempts
-
-    assert {
-        "result": True,
-        "user": USER_LIST[1],
-        "password": "",
-        "lm_hash": "",
-        "ntlm_hash": NT_HASH_LIST[1],
-        "ssh_key": "",
-    } in powershell_exploiter.exploit_attempts
+    # Total 6 attempts reported, 5 failed and 1 succeeded
+    assert len(powershell_exploiter.exploit_attempts) == len(execute_cmd_returns)
+    assert (
+        len([attempt for attempt in powershell_exploiter.exploit_attempts if not attempt["result"]])
+        == 5
+    )
+    assert (
+        len([attempt for attempt in powershell_exploiter.exploit_attempts if attempt["result"]])
+        == 1
+    )
 
 
 def test_build_monkey_execution_command():