diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py index daeb9ea5b..1a0e872f6 100644 --- a/monkey/monkey_island/cc/models/monkey.py +++ b/monkey/monkey_island/cc/models/monkey.py @@ -123,7 +123,8 @@ class Monkey(Document): self.save() -# Can't make following methods static under Monkey class due to ring bug +# TODO Can't make following methods static under Monkey class due to ring bug. When ring will support static methods, we +# should move to static methods in the Monkey class. @ring.lru( expire=1 # data has TTL of 1 second. This is useful for rapid calls for report generation. ) diff --git a/monkey/monkey_island/cc/models/test_monkey.py b/monkey/monkey_island/cc/models/test_monkey.py index fb9b329b1..472c5770b 100644 --- a/monkey/monkey_island/cc/models/test_monkey.py +++ b/monkey/monkey_island/cc/models/test_monkey.py @@ -1,11 +1,15 @@ import uuid +import logging from time import sleep -from .monkey import Monkey -from monkey_island.cc.models.monkey import MonkeyNotFoundError, is_monkey, get_monkey_label_by_id +import pytest + +from monkey_island.cc.models.monkey import Monkey, MonkeyNotFoundError, is_monkey, get_monkey_label_by_id from monkey_island.cc.testing.IslandTestCase import IslandTestCase from .monkey_ttl import MonkeyTtl +logger = logging.getLogger(__name__) + class TestMonkey(IslandTestCase): """ @@ -32,7 +36,7 @@ class TestMonkey(IslandTestCase): # MIA stands for Missing In Action mia_monkey_ttl = MonkeyTtl.create_ttl_expire_in(30) mia_monkey_ttl.save() - mia_monkey = Monkey(guid=str(uuid.uuid4()), dead=False, ttl_ref=mia_monkey_ttl) + mia_monkey = Monkey(guid=str(uuid.uuid4()), dead=False, ttl_ref=mia_monkey_ttl.id) mia_monkey.save() # Emulate timeout - ttl is manually deleted here, since we're using mongomock and not a real mongo instance. sleep(1) @@ -70,8 +74,10 @@ class TestMonkey(IslandTestCase): # Act + assert # Find the existing one self.assertIsNotNone(Monkey.get_single_monkey_by_id(a_monkey.id)) + # Raise on non-existent monkey - self.assertRaises(MonkeyNotFoundError, Monkey.get_single_monkey_by_id, "abcdefabcdefabcdefabcdef") + with pytest.raises(MonkeyNotFoundError) as e_info: + _ = Monkey.get_single_monkey_by_id("abcdefabcdefabcdefabcdef") def test_get_os(self): self.fail_if_not_testing_env() @@ -125,29 +131,41 @@ class TestMonkey(IslandTestCase): ip_addresses=[ip_example]) linux_monkey.save() + logger.debug(id(get_monkey_label_by_id)) + cache_info_before_query = get_monkey_label_by_id.storage.backend.cache_info() self.assertEqual(cache_info_before_query.hits, 0) + self.assertEqual(cache_info_before_query.misses, 0) # not cached label = get_monkey_label_by_id(linux_monkey.id) + cache_info_after_query_1 = get_monkey_label_by_id.storage.backend.cache_info() + self.assertEqual(cache_info_after_query_1.hits, 0) + self.assertEqual(cache_info_after_query_1.misses, 1) + logger.info("1) ID: {} label: {}".format(linux_monkey.id, label)) self.assertIsNotNone(label) self.assertIn(hostname_example, label) self.assertIn(ip_example, label) # should be cached - _ = get_monkey_label_by_id(linux_monkey.id) - cache_info_after_query = get_monkey_label_by_id.storage.backend.cache_info() - self.assertEqual(cache_info_after_query.hits, 1) + label = get_monkey_label_by_id(linux_monkey.id) + logger.info("2) ID: {} label: {}".format(linux_monkey.id, label)) + cache_info_after_query_2 = get_monkey_label_by_id.storage.backend.cache_info() + self.assertEqual(cache_info_after_query_2.hits, 1) + self.assertEqual(cache_info_after_query_2.misses, 1) + # set hostname deletes the id from the cache. linux_monkey.set_hostname("Another hostname") # should be a miss label = get_monkey_label_by_id(linux_monkey.id) - cache_info_after_second_query = get_monkey_label_by_id.storage.backend.cache_info() + logger.info("3) ID: {} label: {}".format(linux_monkey.id, label)) + cache_info_after_query_3 = get_monkey_label_by_id.storage.backend.cache_info() + logger.debug("Cache info: {}".format(str(cache_info_after_query_3))) # still 1 hit only - self.assertEqual(cache_info_after_second_query.hits, 1) - self.assertEqual(cache_info_after_second_query.misses, 2) + self.assertEqual(cache_info_after_query_3.hits, 1) + self.assertEqual(cache_info_after_query_3.misses, 2) def test_is_monkey(self): self.fail_if_not_testing_env()