working on tests for disk_storage

shoebox/disk_storage.py

@@ -3,22 +3,21 @@ import datetime
 import json
 import struct
 
-def handle_datetime(obj):
-    if isinstance(obj, datetime.datetime):
-        if obj.utcoffset() is not None:
-            obj = obj - obj.utcoffset()
-    millis = int(
-        calendar.timegm(obj.timetuple()) * 1000 +
-        obj.microsecond / 1000
-    )
-    return millis
+class DatetimeEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, datetime.datetime):
+            if obj.utcoffset() is not None:
+                obj = obj - obj.utcoffset()
+            return int(calendar.timegm(obj.timetuple()) * 1000 +
+                         obj.microsecond / 1000)
+        return super(DatetimeEncoder, self).default(obj)
 
 
-def InvalidVersion(Exception):
+class InvalidVersion(Exception):
     pass
 
 
-def OutOfSync(Exception):
+class OutOfSync(Exception):
     pass
 
 
@@ -40,8 +39,9 @@ class Version0(object):
     def load_preamble(self, file_handle):
         raw = file_handle.read(self.preamble_size)
         header = struct.unpack(self.preamble_schema, raw)
+        print "raw", raw
         if header[0] != BOR_MAGIC_NUMBER:
-            raise OutOfSync()
+            raise OutOfSync("Expected Beginning of Record marker")
         return header[1]
 
 
@@ -180,7 +180,7 @@ if __name__ == "__main__":
              }
             }
 
-    json_event = json.dumps(event, default=handle_datetime)
+    json_event = json.dumps(event, cls=DatetimeEncoder)
     metadata = {'request_id': event['request_id'],
                 'event_type': event['event_type'],
                 'source': event['source'],

test/test_disk_storage.py

@@ -0,0 +1,122 @@
+import datetime
+import mock
+import json
+import struct
+import unittest
+
+import dateutil.tz
+
+from shoebox import disk_storage
+
+
+class TestDiskStorage(unittest.TestCase):
+    def setUp(self):
+        self.handler = disk_storage.DatetimeEncoder()
+
+    def test_handle_datetime_non_datetime(self):
+        self.assertRaises(TypeError, self.handler.default, "text")
+
+    def test_handle_datetime(self):
+        now = datetime.datetime(day=1, month=2, year=2014,
+                                hour=10, minute=11, second=12)
+        self.assertEqual(1391249472000, self.handler.default(now))
+
+    def test_handle_datetime_offset(self):
+        now = datetime.datetime(day=1, month=2, year=2014,
+                                hour=10, minute=11, second=12,
+                                tzinfo=dateutil.tz.tzoffset(None, 4*60*60))
+        self.assertEqual(1391235072000, self.handler.default(now))
+
+
+class TestVersion0(unittest.TestCase):
+    def setUp(self):
+       self.v0 = disk_storage.Version0()
+
+    def test_make_preamble(self):
+       self.assertEqual(6, len(self.v0.make_preamble(99)))
+
+    def test_load_preamble_bad_bor(self):
+        file_handle = mock.Mock()
+        file_handle.read.return_value = "abcdef"
+        self.assertRaises(disk_storage.OutOfSync, self.v0.load_preamble,
+                          file_handle)
+
+    def test_load_preamble(self):
+        file_handle = mock.Mock()
+        file_handle.read.return_value = struct.pack("ih",
+                          disk_storage.BOR_MAGIC_NUMBER, 99)
+        self.assertEqual(99, self.v0.load_preamble(file_handle))
+
+
+class TestVersion1(unittest.TestCase):
+    def setUp(self):
+       self.v1 = disk_storage.Version1()
+
+    def test_no_metadata(self):
+        metadata = {}
+        payload = "shoebox"
+        package = self.v1.pack(payload, metadata)
+        self.assertEqual(4, len(package))
+        self.assertEqual(12, len(package[1]))  # header
+        self.assertEqual(4, len(package[2]))  # metadata
+        self.assertEqual("\x00\x00\x00\x00", package[2])
+        self.assertEqual(11, len(package[3]))  # payload 4+7
+
+    def test_empty_payload(self):
+        metadata = {"key": "value", "some": "stuff"}
+        payload = ""
+        package = self.v1.pack(payload, metadata)
+        self.assertEqual(4, len(package))
+        self.assertEqual(12, len(package[1]))  # header
+        self.assertEqual(37, len(package[2]))  # metadata 4+(4*4)+3+5+4+5
+        self.assertEqual(4, len(package[3]))  # payload 4+0
+        self.assertEqual("\x00\x00\x00\x00", package[3])
+
+    def test_unpack_happy_day(self):
+        metadata = {"key": "value", "some": "stuff"}
+        payload = {"shoebox": 1234}
+        jpayload = json.dumps(payload)
+        blocks = self.v1.pack(jpayload, metadata)
+        blocks = blocks[1:]  # Remove preamble
+
+        file_handle = mock.Mock()
+        file_handle.read.side_effect = blocks
+
+        m, p = self.v1.unpack(file_handle)
+        self.assertEqual(metadata, m)
+        self.assertEqual(payload, p)
+
+    def test_unpack_bad_eor(self):
+        metadata = {"key": "value", "some": "stuff"}
+        payload = {"shoebox": 1234}
+        jpayload = json.dumps(payload)
+        blocks = self.v1.pack(jpayload, metadata)
+        blocks = blocks[1:]  # Remove preamble
+
+        # break the EOR marker
+        print len(blocks[0])
+        newblock = blocks[0][:8] + '\x00\x00\x01\x02'
+        blocks = list(blocks)
+        blocks[0] = newblock
+        blocks = tuple(blocks)
+
+        file_handle = mock.Mock()
+        file_handle.read.side_effect = blocks
+
+        self.assertRaises(disk_storage.OutOfSync, self.v1.unpack, file_handle)
+
+
+class TestHelpers(unittest.TestCase):
+    def test_get_version_handler_bad(self):
+        self.assertRaises(disk_storage.InvalidVersion,
+                          disk_storage.get_version_handler, 99)
+
+    def test_get_version_handler(self):
+        self.assertTrue(isinstance(disk_storage.get_version_handler(1),
+                                   disk_storage.Version1))
+
+        # Default version ...
+        self.assertTrue(isinstance(disk_storage.get_version_handler(),
+                                   disk_storage.Version1))
+
+

test/test_roll_checker.py

@@ -0,0 +1,50 @@
+import datetime
+import mock
+import unittest
+
+from shoebox import roll_checker
+from shoebox import utils
+
+
+class TestRollChecker(unittest.TestCase):
+    def test_time_roll_checker_start(self):
+        one_hour = datetime.timedelta(hours=1)
+        x = roll_checker.TimeRollChecker(one_hour)
+        now = datetime.datetime.utcnow()
+        with mock.patch.object(utils, 'now') as dt:
+            dt.return_value = now
+            x.start(None)
+        self.assertEqual(x.start_time, now)
+        self.assertEqual(x.end_time, now + one_hour)
+
+    def test_time_roll_checker_end(self):
+        one_hour = datetime.timedelta(hours=1)
+        x = roll_checker.TimeRollChecker(one_hour)
+        now = datetime.datetime.utcnow()
+        x.start_time = now
+        x.end_time = now + one_hour
+        with mock.patch.object(utils, 'now') as dt:
+            dt.return_value = now + one_hour
+            self.assertTrue(x.check(None))
+
+        with mock.patch.object(utils, 'now') as dt:
+            dt.return_value = now
+            self.assertFalse(x.check(None))
+
+        with mock.patch.object(utils, 'now') as dt:
+            dt.return_value = now + one_hour - datetime.timedelta(seconds = 1)
+            self.assertFalse(x.check(None))
+
+    def test_size_roll_checker_end(self):
+        one_gig = 1073741824
+        x = roll_checker.SizeRollChecker(10)
+
+        archive = mock.Mock()
+        archive._get_file_handle.return_value.tell.return_value = one_gig * 5
+        self.assertFalse(x.check(archive))
+
+        archive._get_file_handle.return_value.tell.return_value = one_gig * 10
+        self.assertTrue(x.check(archive))
+
+        archive._get_file_handle.return_value.tell.return_value = one_gig * 11
+        self.assertTrue(x.check(archive))

test/test_roll_manager.py

@@ -0,0 +1,73 @@
+import datetime
+import mock
+import os
+import unittest
+
+from shoebox import archive
+from shoebox import roll_checker
+from shoebox import roll_manager
+from shoebox import utils
+
+
+class TestRollManager(unittest.TestCase):
+    def test_make_filename(self):
+        now = datetime.datetime(day=1, month=2, year=2014,
+                                hour=10, minute=11, second=12)
+        x = roll_manager.RollManager("filename_%c.dat", None)
+
+        with mock.patch.object(utils, "now") as dt:
+            dt.return_value = now
+            filename = x._make_filename()
+            self.assertEqual(filename, "filename_Sat_Feb__1_10:11:12_2014.dat")
+
+
+class FakeArchive(object):
+    def __init__(self, filename):
+        self.filename = filename
+        self.data = []
+
+    def write(self, metadata, payload):
+        self.data.append((metadata, payload))
+
+
+class TestWritingRollManager(unittest.TestCase):
+    def test_get_active_archive(self):
+        checker = mock.Mock()
+        filename_template = "filename_%c.dat"
+        x = roll_manager.WritingRollManager(filename_template, checker)
+        with mock.patch("shoebox.archive.ArchiveWriter._open_file") as of:
+            arc = x.get_active_archive()
+            self.assertTrue(isinstance(arc, archive.ArchiveWriter))
+            self.assertTrue(checker.start.called)
+
+    def test_write_always_roll(self):
+        checker = mock.Mock()
+        checker.check.return_value = True
+        x = roll_manager.WritingRollManager("template", checker,
+                                            archive_class=FakeArchive)
+        with mock.patch.object(x, "_roll_archive") as ra:
+            x.write({}, "payload")
+            self.assertTrue(ra.called)
+
+    def test_write_never_roll(self):
+        checker = mock.Mock()
+        checker.check.return_value = False
+        x = roll_manager.WritingRollManager("template", checker,
+                                            archive_class=FakeArchive)
+        with mock.patch.object(x, "_roll_archive") as ra:
+            x.write({}, "payload")
+            self.assertFalse(ra.called)
+
+
+class TestWriting(unittest.TestCase):
+    def test_write(self):
+        checker = roll_checker.NeverRollChecker()
+        x = roll_manager.WritingRollManager("template_%s", checker,
+                                            archive_class=FakeArchive)
+
+        for index in range(10):
+            x.write({"index": str(index)}, "payload_%d" % index)
+
+        arc = x.get_active_archive()
+        self.assertEqual(10, len(arc.data))
+        self.assertEqual(({"index": "0"}, "payload_0"), arc.data[0])

tox.ini

@@ -3,9 +3,9 @@ envlist = py26,py27
 
 [testenv]
 deps = 
+    python-dateutil
     nose
     mock
     coverage
-    pymongo
 
 commands = nosetests -d --with-coverage --cover-inclusive --cover-package shoebox []