import datetime import mock import os import unittest import notification_utils from shoebox import archive from shoebox import roll_checker from shoebox import roll_manager class FakeArchive(object): def __init__(self, filename): self.filename = filename self.data = [] def write(self, metadata, payload): self.data.append((metadata, payload)) class TestRollManager(unittest.TestCase): def test_close(self): callback = mock.Mock() x = roll_manager.RollManager("template", archive_callback=callback) x.active_archive = mock.Mock() x.active_filename = "foo" x.close() self.assertTrue(x.active_archive is None) self.assertTrue(x.active_filename is None) self.assertTrue(callback.on_close.called) class TestWritingRollManager(unittest.TestCase): 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) def test_correct_archiver(self): x = roll_manager.WritingRollManager("foo", None) print x.archive_class self.assertEqual(x.archive_class, archive.ArchiveWriter) def test_get_active_archive(self): checker = mock.Mock() callback = mock.Mock() filename_template = "filename_%c.dat" x = roll_manager.WritingRollManager(filename_template, checker, archive_callback=callback, archive_class=FakeArchive) with mock.patch("shoebox.archive.ArchiveWriter._open_file") as of: arc = x.get_active_archive() self.assertTrue(checker.start.called) self.assertTrue(callback.on_open.called) def test_make_filename(self): now = datetime.datetime(day=1, month=2, year=2014, hour=10, minute=11, second=12) x = roll_manager.WritingRollManager("filename_%c.dat", None) with mock.patch.object(notification_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 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]) class TestJSONRollManager(unittest.TestCase): def test_make_filename(self): now = datetime.datetime(day=1, month=2, year=2014, hour=10, minute=11, second=12) with mock.patch.object(notification_utils, "now") as dt: with mock.patch.object(notification_utils, "dt_to_decimal") as td: td.return_value = 123.45 dt.return_value = now x = roll_manager.WritingJSONRollManager( "%Y%m%d [[TIMESTAMP]] [[CRC]].foo") fn = x._make_filename("mycrc", "foo") self.assertEqual("foo/20140201_123.45_mycrc.foo", fn) @mock.patch('os.path.getsize') @mock.patch('os.listdir') @mock.patch('os.path.isfile') def test_should_tar(self, isf, ld, gs): rm = roll_manager.WritingJSONRollManager("template.foo") gs.return_value = 250000 ld.return_value = ['a', 'b', 'c'] isf.return_value = True rm.roll_size_mb = 1 self.assertFalse(rm._should_tar()) ld.return_value = ['a', 'b', 'c', 'd', 'e', 'f'] self.assertTrue(rm._should_tar()) @mock.patch('os.listdir') @mock.patch('os.remove') @mock.patch('os.path.isfile') def test_clean_working_directory(self, isf, rem, ld): isf.return_value = True rm = roll_manager.WritingJSONRollManager("template.foo") ld.return_value = ['a', 'b', 'c'] rm._clean_working_directory() self.assertEqual(3, rem.call_count) @mock.patch('os.listdir') @mock.patch('tarfile.open') @mock.patch('os.path.isfile') def test_tar_directory(self, isf, to, ld): ld.return_value = ['a', 'b', 'c'] isf.return_value = True rm = roll_manager.WritingJSONRollManager("template.foo") open_name = '%s.open' % roll_manager.__name__ with mock.patch(open_name, create=True) as mock_open: mock_open.return_value = mock.MagicMock(spec=file) rm._tar_directory() self.assertTrue(to.called) def test_write(self): rm = roll_manager.WritingJSONRollManager("template.foo") payload = "some big payload" open_name = '%s.open' % roll_manager.__name__ with mock.patch(open_name, create=True) as mock_open: with mock.patch.object(rm, "_should_tar") as st: with mock.patch.object(rm, "_tar_directory") as td: st.return_value = False mock_open.return_value = mock.MagicMock(spec=file) rm.write("metadata", payload) self.assertTrue(mock_open.called_once_with( "template.foo", "wb")) self.assertFalse(td.called)