commit 9b06da8803e989b70210f27ac97226abf2f3c494 Author: Arturo Filastò <arturo@filasto.net> Date: Tue Oct 11 18:02:23 2016 +0200 Add two new unittests that ensure the scheduler workflow is working * Also ensures that a deck is run twice by the scheduler --- ooni/tests/test_scheduler.py | 260 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 259 insertions(+), 1 deletion(-) diff --git a/ooni/tests/test_scheduler.py b/ooni/tests/test_scheduler.py index c8eaf55..e954668 100644 --- a/ooni/tests/test_scheduler.py +++ b/ooni/tests/test_scheduler.py @@ -1,13 +1,24 @@ import os import sys +import json import shutil import random import tempfile +import mock + +from datetime import datetime, timedelta + from twisted.internet import reactor, task, defer from twisted.trial import unittest -from ooni.agent.scheduler import ScheduledTask, DidNotRun, FileSystemlockAndMutex +from ooni.utils import generate_filename, randomSTR, randomDate, LONG_DATE +from ooni.tests.bases import ConfigTestCase + +from ooni.deck.store import DeckStore +from ooni.agent.scheduler import ScheduledTask, DidNotRun +from ooni.agent.scheduler import FileSystemlockAndMutex +from ooni.agent.scheduler import SchedulerService class TestScheduler(unittest.TestCase): def test_scheduled_task(self): @@ -112,3 +123,250 @@ class TestScheduler(unittest.TestCase): self.assertFalse(lock.locked) shutil.rmtree(lock_dir) + +def random_measurement_name(start_date=None, end_date=None): + # By default we use as start date something in the past 6 days and end + # date today. + if start_date is None: + start_date = datetime.now() - timedelta(days=6) + if end_date is None: + end_date = datetime.now() + + test_details = dict( + test_name=random.choice(['http_invalid_request_line', + 'web_connectivity', + 'http_header_field_manipulation', + 'vanilla_tor', + 'new_test_name']), + probe_cc=randomSTR(2, num=False), # XXX this should be a valid CC + probe_asn='AS'+str(random.randint(0, 2**16)), + test_start_time=randomDate(start_date, end_date).strftime(LONG_DATE) + ) + return generate_filename(test_details) + +def get_measurement_header(): + return { + 'probe_asn': 'AS'+str(random.randint(0, 2**16)), + 'probe_cc': randomSTR(2, num=False), + 'probe_ip': '127.0.0.1', + 'probe_city': None, + 'software_name': 'ooniprobe', + 'software_version': '0.0.0', + 'options': {}, + 'annotations': {}, + 'data_format_version': '0.2.0', + 'test_name': 'dummy', + 'test_version': '0.0.0', + 'test_helpers': {}, + 'test_start_time': '2016-01-01 01:01:01', + 'test_runtime': 0.1, + 'input_hashes': [], + 'report_id': randomSTR(100), + 'test_keys': {}, + 'input': '' + } + +def write_dummy_measurements(fh, size=100): + """ + :param fh: an open file handle + :param size: size of the measurements in bytes to write + :return: The actual size that has been written. + """ + written_size = 0 + while written_size < size: + entry = get_measurement_header() + entry['test_keys']['data'] = randomSTR(int(size / 10)) + data = json.dumps(entry) + written_size += len(data) + fh.write(data) + return written_size + + +DUMMY_WEB_FULL = """ +--- +name: Full Web test deck +description: This deck runs HTTP Header Field Manipulation, HTTP Invalid + Request and the Web Connectivity test +schedule: "@daily" +tasks: +- name: Runs the HTTP Header Field Manipulation test + ooni: + test_name: http_header_field_manipulation +""" + +class TestSchedulerService(ConfigTestCase): + def setUp(self): + super(TestSchedulerService, self).setUp() + + self.config_patcher = mock.patch('ooni.agent.scheduler.config') + self.config_m_patcher = mock.patch('ooni.measurements.config') + self.config_mock = self.config_patcher.start() + self.config_m_mock = self.config_m_patcher.start() + + self.scheduler_directory = tempfile.mkdtemp() + self.config_mock.scheduler_directory = self.scheduler_directory + + self.running_path = tempfile.mkdtemp() + self.config_mock.running_path = self.running_path + + self.config_mock.is_initialized.return_value = True + + self.config_mock.basic.measurement_quota = '100M' + + self.measurements_directory = tempfile.mkdtemp() + self.config_mock.measurements_directory = self.measurements_directory + self.create_dummy_measurements() + + self.config_m_mock.measurements_directory = self.measurements_directory + + self.decks_enabled_directory = tempfile.mkdtemp() + self.decks_available_directory = tempfile.mkdtemp() + + self.deck_store = DeckStore( + available_directory=self.decks_available_directory, + enabled_directory=self.decks_enabled_directory + ) + + self.mock_deck = mock.MagicMock() + self.deck_store.get = lambda deck_id: self.mock_deck + self.mock_deck.setup.return_value = defer.succeed(None) + self.mock_deck.run.return_value = defer.succeed(None) + + with open(os.path.join(self.decks_available_directory, + 'web-full.yaml'), 'w') as out_file: + out_file.write(DUMMY_WEB_FULL) + + def create_dummy_measurements(self, count=10, size=10*1024): + for _ in range(count): + dir_path = os.path.join( + self.measurements_directory, + random_measurement_name() + ) + os.mkdir(dir_path) + with open(os.path.join(dir_path, "measurements.njson"), 'w') as fh: + write_dummy_measurements(fh, float(size) / float(count)) + + def tearDown(self): + super(TestSchedulerService, self).tearDown() + + shutil.rmtree(self.measurements_directory) + shutil.rmtree(self.scheduler_directory) + shutil.rmtree(self.running_path) + + self.config_patcher.stop() + self.config_m_patcher.stop() + + @mock.patch('ooni.agent.scheduler.resources') + @mock.patch('ooni.agent.scheduler.probe_ip') + @mock.patch('ooni.agent.scheduler.input_store') + @mock.patch('ooni.agent.scheduler.oonireport') + def test_deck_run_twice(self, mock_resources, mock_probe_ip, + mock_input_store, mock_oonireport): + mock_probe_ip.geodata['countrycode'] = 'ZZ' + mock_probe_ip.lookup.return_value = defer.succeed(None) + mock_probe_ip.resolveGeodata.return_value = defer.succeed(None) + + mock_resources.check_for_update.return_value = defer.succeed(None) + + mock_input_store.update.return_value = defer.succeed(None) + + mock_oonireport.upload_all.return_value = defer.succeed(None) + + mock_director = mock.MagicMock() + d = defer.Deferred() + with mock.patch('ooni.agent.scheduler.deck_store', self.deck_store): + + dummy_clock = task.Clock() + scheduler_service = SchedulerService( + director=mock_director, + _reactor=dummy_clock + ) + scheduler_service.startService() + dummy_clock.advance(30) + + now_time = datetime.utcnow() + DT_FRMT = "%Y-%m-%dT%H:%M:%SZ" + + for t in scheduler_service._scheduled_tasks: + with open(os.path.join(self.scheduler_directory, + t.identifier)) as in_file: + dstr = datetime.strptime(in_file.read(), + DT_FRMT).strftime("%Y-%m-%dT%H") + self.assertEqual(dstr, now_time.strftime("%Y-%m-%dT%H")) + + dummy_clock.advance(30) + dummy_clock.advance(30) + dummy_clock.advance(30) + dummy_clock.advance(30) + dummy_clock.advance(30) + dummy_clock.advance(30) + # Here we pretend they ran yesterday so to re-trigger the daily + # tasks + for t in scheduler_service._scheduled_tasks: + with open(os.path.join(self.scheduler_directory, + t.identifier), 'w') as out_file: + yesterday = (now_time - timedelta(days=1, + hours=2)).strftime(DT_FRMT) + out_file.write(yesterday) + dummy_clock.advance(30) + + # We check that the run method of the deck was called twice + self.mock_deck.run.assert_has_calls([ + mock.call(mock_director), mock.call(mock_director) + ]) + d.callback(None) + + return d + + @mock.patch('ooni.agent.scheduler.resources') + @mock.patch('ooni.agent.scheduler.probe_ip') + @mock.patch('ooni.agent.scheduler.input_store') + @mock.patch('ooni.agent.scheduler.oonireport') + def test_disk_quota_cleanup(self, mock_resources, mock_probe_ip, + mock_input_store, mock_oonireport): + + mock_probe_ip.geodata['countrycode'] = 'ZZ' + mock_probe_ip.lookup.return_value = defer.succeed(None) + mock_probe_ip.resolveGeodata.return_value = defer.succeed(None) + + mock_resources.check_for_update.return_value = defer.succeed(None) + + mock_input_store.update.return_value = defer.succeed(None) + + mock_oonireport.upload_all.return_value = defer.succeed(None) + + self.config_mock.basic.measurement_quota = '1M' + # We create 10MB of measurements + self.create_dummy_measurements(count=10, size=1*1024*1024) + measurement_count = len(os.listdir(self.measurements_directory)) + + mock_director = mock.MagicMock() + d = defer.Deferred() + with mock.patch('ooni.agent.scheduler.deck_store', self.deck_store): + + dummy_clock = task.Clock() + scheduler_service = SchedulerService( + director=mock_director, + _reactor=dummy_clock + ) + scheduler_service.startService() + dummy_clock.advance(30) + + now_time = datetime.utcnow() + DT_FRMT = "%Y-%m-%dT%H:%M:%SZ" + + for t in scheduler_service._scheduled_tasks: + with open(os.path.join(self.scheduler_directory, + t.identifier)) as in_file: + dstr = datetime.strptime(in_file.read(), + DT_FRMT).strftime("%Y-%m-%dT%H") + self.assertEqual(dstr, now_time.strftime("%Y-%m-%dT%H")) + + # Ensure there are less measurements than there were at the + # beginning + new_measurement_count = len(os.listdir(self.measurements_directory)) + self.assertGreater(measurement_count, new_measurement_count) + + d.callback(None) + + return d