From 411b027e923191b12f442d9bacfbbeffb75f0016 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 30 Sep 2022 10:15:31 -0400 Subject: [PATCH] Island: Implement MongoAgentRepository.get_agents() --- .../cc/repository/mongo_agent_repository.py | 7 ++++++- .../repository/test_mongo_agent_repository.py | 20 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/repository/mongo_agent_repository.py b/monkey/monkey_island/cc/repository/mongo_agent_repository.py index 70aa0a2ed..9bf66e4bc 100644 --- a/monkey/monkey_island/cc/repository/mongo_agent_repository.py +++ b/monkey/monkey_island/cc/repository/mongo_agent_repository.py @@ -39,7 +39,12 @@ class MongoAgentRepository(IAgentRepository): ) def get_agents(self) -> Sequence[Agent]: - return [] + try: + cursor = self._agents_collection.find({}, {MONGO_OBJECT_ID_KEY: False}) + except Exception as err: + raise RetrievalError(f"Error retrieving agents: {err}") + + return [Agent(**a) for a in cursor] def get_agent_by_id(self, agent_id: AgentID) -> Agent: try: diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py index 62ebaed1c..51999d1c3 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py @@ -161,6 +161,26 @@ def test_upsert_agent__storage_error_insert_failed(error_raising_mock_mongo_clie agent_repository.upsert_agent(agent) +def test_get_agents__empty_repo(empty_agent_repository): + all_agents = empty_agent_repository.get_agents() + + assert len(all_agents) == 0 + + +def test_get_agents(agent_repository): + all_agents = agent_repository.get_agents() + + assert len(all_agents) == len(AGENTS) + + for agent in AGENTS: + assert agent in all_agents + + +def test_get_agents__retrieval_error(error_raising_agent_repository): + with pytest.raises(RetrievalError): + error_raising_agent_repository.get_agents() + + def test_get_agent_by_id(agent_repository): for i, expected_agent in enumerate(AGENTS): assert agent_repository.get_agent_by_id(expected_agent.id) == expected_agent