From 5c97d6bdbf035f2b568050c67b421c01d39f8e62 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 21 Jun 2022 11:00:13 -0400 Subject: [PATCH] UT: Add tests for AgentBinaryRepository --- .../test_agent_binary_repository.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_binary_repository.py diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_binary_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_binary_repository.py new file mode 100644 index 000000000..585867b2e --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_agent_binary_repository.py @@ -0,0 +1,31 @@ +from unittest.mock import MagicMock + +import pytest + +from monkey_island.cc.repository import AgentBinaryRepository, IFileRepository, RetrievalError + +LINUX_AGENT_BINARY = b"linux_binary" +WINDOWS_AGENT_BINARY = b"windows_binary" + + +@pytest.fixture +def error_mock_file_repository(): + mock_file_repository = MagicMock(wraps=IFileRepository) + mock_file_repository.open_file = MagicMock(side_effect=OSError) + + return mock_file_repository + + +@pytest.fixture +def agent_binary_repository(error_mock_file_repository): + return AgentBinaryRepository(error_mock_file_repository) + + +def test_get_linux_binary_retrieval_error(agent_binary_repository): + with pytest.raises(RetrievalError): + agent_binary_repository.get_linux_binary() + + +def test_get_windows_binary_retrieval_error(agent_binary_repository): + with pytest.raises(RetrievalError): + agent_binary_repository.get_windows_binary()