diff --git a/django/website/data_layer/handlers.py b/django/website/data_layer/handlers.py index df9c23be3926525bea84c3f69302366753849f0d..2d31bff11699106826ea62b658e067f1f0a641f3 100644 --- a/django/website/data_layer/handlers.py +++ b/django/website/data_layer/handlers.py @@ -6,3 +6,8 @@ class Message(object): @classmethod def create(cls, message): models.Message(**message).save() + + @classmethod + def list(cls): + # TODO: I think we should probably use a DRF serializer here + return models.Message.objects.values().iterator() diff --git a/django/website/data_layer/tests/handler_tests.py b/django/website/data_layer/tests/handler_tests.py index 52e264eb786ad7bd3548cfe46d62c35110a8e3d9..99345a651a6110ae2bd6978a7a7e2b707e789882 100644 --- a/django/website/data_layer/tests/handler_tests.py +++ b/django/website/data_layer/tests/handler_tests.py @@ -20,3 +20,15 @@ class HandlerTests(TestCase): outmessage = models.Message.objects.get() self.assertEqual(outmessage.body, "Text") self.assertEqual(outmessage.timestamp, now) + + def test_message_list_empty(self): + messages = handlers.Message.list() + self.assertEqual(len(list(messages)), 0) + + def test_message_list_one(self): + models.Message(body="Test").save() + + [message] = handlers.Message.list() + + self.assertEqual(message['body'], "Test") +