Added get_single_monkey_by_id API to the monkey model

Useful in many cases in the code.
Also added unittest for this method which passed.
This commit is contained in:
Shay Nehmad 2019-05-07 15:56:27 +03:00
parent 295525dfed
commit 3c6cda03af
3 changed files with 18 additions and 1 deletions
monkey/monkey_island/cc/models

View File

@ -0,0 +1,2 @@
class MonkeyNotFoundError(Exception):
pass

View File

@ -5,6 +5,7 @@ import mongoengine
from mongoengine import Document, StringField, ListField, BooleanField, EmbeddedDocumentField, DateField, \
ReferenceField
from monkey_island.cc.models.errors import MonkeyNotFoundError
from monkey_island.cc.models.monkey_ttl import MonkeyTtl
@ -32,6 +33,13 @@ class Monkey(Document):
pba_results = ListField()
ttl_ref = ReferenceField(MonkeyTtl)
@staticmethod
def get_single_monkey_by_id(db_id):
try:
return Monkey.objects(id=db_id)[0]
except IndexError:
raise MonkeyNotFoundError("id: {0}".format(str(db_id)))
def is_dead(self):
monkey_is_dead = False
if self.dead:
@ -41,7 +49,7 @@ class Monkey(Document):
if MonkeyTtl.objects(id=self.ttl_ref.id).count() == 0:
# No TTLs - monkey has timed out. The monkey is MIA.
monkey_is_dead = True
except mongoengine.DoesNotExist:
except (mongoengine.DoesNotExist, AttributeError):
# Trying to dereference unknown document - the monkey is MIA.
monkey_is_dead = True
return monkey_is_dead

View File

@ -7,6 +7,7 @@ from unittest import TestCase
import mongomock
from monkey import Monkey
from monkey_island.cc.models.errors import MonkeyNotFoundError
from monkey_ttl import MonkeyTtl
@ -35,3 +36,9 @@ class TestMonkey(TestCase):
self.assertTrue(mia_monkey.is_dead())
self.assertFalse(alive_monkey.is_dead())
def test_get_single_monkey_by_id(self):
a_monkey = Monkey(guid=str(uuid.uuid4()))
a_monkey.save()
self.assertIsNotNone(Monkey.get_single_monkey_by_id(a_monkey.id))
self.assertRaises(MonkeyNotFoundError, Monkey.get_single_monkey_by_id, "abcdefabcdefabcdefabcdef")