diff --git a/shoebox/disk_storage.py b/shoebox/disk_storage.py index 9001620..98924e3 100644 --- a/shoebox/disk_storage.py +++ b/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'], diff --git a/test/test_disk_storage.py b/test/test_disk_storage.py new file mode 100644 index 0000000..6069c42 --- /dev/null +++ b/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)) + + diff --git a/test/test_roll_checker.py b/test/test_roll_checker.py new file mode 100644 index 0000000..9e911a0 --- /dev/null +++ b/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)) diff --git a/test/test_roll_manager.py b/test/test_roll_manager.py new file mode 100644 index 0000000..d2c9409 --- /dev/null +++ b/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]) diff --git a/tox.ini b/tox.ini index 50f5ea1..9118d5a 100644 --- a/tox.ini +++ b/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 []