Fix TestMonkey to pytest

This commit is contained in:
Shay Nehmad 2019-10-27 17:55:51 +02:00
parent 7b153d29b2
commit f2297de661
2 changed files with 30 additions and 11 deletions

View File

@ -123,7 +123,8 @@ class Monkey(Document):
self.save() 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( @ring.lru(
expire=1 # data has TTL of 1 second. This is useful for rapid calls for report generation. expire=1 # data has TTL of 1 second. This is useful for rapid calls for report generation.
) )

View File

@ -1,11 +1,15 @@
import uuid import uuid
import logging
from time import sleep from time import sleep
from .monkey import Monkey import pytest
from monkey_island.cc.models.monkey import MonkeyNotFoundError, is_monkey, get_monkey_label_by_id
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_island.cc.testing.IslandTestCase import IslandTestCase
from .monkey_ttl import MonkeyTtl from .monkey_ttl import MonkeyTtl
logger = logging.getLogger(__name__)
class TestMonkey(IslandTestCase): class TestMonkey(IslandTestCase):
""" """
@ -32,7 +36,7 @@ class TestMonkey(IslandTestCase):
# MIA stands for Missing In Action # MIA stands for Missing In Action
mia_monkey_ttl = MonkeyTtl.create_ttl_expire_in(30) mia_monkey_ttl = MonkeyTtl.create_ttl_expire_in(30)
mia_monkey_ttl.save() 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() mia_monkey.save()
# Emulate timeout - ttl is manually deleted here, since we're using mongomock and not a real mongo instance. # Emulate timeout - ttl is manually deleted here, since we're using mongomock and not a real mongo instance.
sleep(1) sleep(1)
@ -70,8 +74,10 @@ class TestMonkey(IslandTestCase):
# Act + assert # Act + assert
# Find the existing one # Find the existing one
self.assertIsNotNone(Monkey.get_single_monkey_by_id(a_monkey.id)) self.assertIsNotNone(Monkey.get_single_monkey_by_id(a_monkey.id))
# Raise on non-existent monkey # 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): def test_get_os(self):
self.fail_if_not_testing_env() self.fail_if_not_testing_env()
@ -125,29 +131,41 @@ class TestMonkey(IslandTestCase):
ip_addresses=[ip_example]) ip_addresses=[ip_example])
linux_monkey.save() linux_monkey.save()
logger.debug(id(get_monkey_label_by_id))
cache_info_before_query = get_monkey_label_by_id.storage.backend.cache_info() 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.hits, 0)
self.assertEqual(cache_info_before_query.misses, 0)
# not cached # not cached
label = get_monkey_label_by_id(linux_monkey.id) 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.assertIsNotNone(label)
self.assertIn(hostname_example, label) self.assertIn(hostname_example, label)
self.assertIn(ip_example, label) self.assertIn(ip_example, label)
# should be cached # should be cached
_ = get_monkey_label_by_id(linux_monkey.id) label = get_monkey_label_by_id(linux_monkey.id)
cache_info_after_query = get_monkey_label_by_id.storage.backend.cache_info() logger.info("2) ID: {} label: {}".format(linux_monkey.id, label))
self.assertEqual(cache_info_after_query.hits, 1) 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") linux_monkey.set_hostname("Another hostname")
# should be a miss # should be a miss
label = get_monkey_label_by_id(linux_monkey.id) 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 # still 1 hit only
self.assertEqual(cache_info_after_second_query.hits, 1) self.assertEqual(cache_info_after_query_3.hits, 1)
self.assertEqual(cache_info_after_second_query.misses, 2) self.assertEqual(cache_info_after_query_3.misses, 2)
def test_is_monkey(self): def test_is_monkey(self):
self.fail_if_not_testing_env() self.fail_if_not_testing_env()