forked from p15670423/monkey
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:
parent
295525dfed
commit
3c6cda03af
monkey/monkey_island/cc/models
|
@ -0,0 +1,2 @@
|
|||
class MonkeyNotFoundError(Exception):
|
||||
pass
|
|
@ -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
|
||||
|
|
|
@ -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")
|
||||
|
|
Loading…
Reference in New Issue