<s> import os <EOL> from twisted . trial import unittest <EOL> from twisted . internet import defer , task <EOL> from ooni . managers import MeasurementManager <EOL> from ooni . tests . mocks import MockSuccessTask , MockFailTask , MockFailOnceTask , MockFailure <EOL> from ooni . tests . mocks import MockSuccessTaskWithTimeout , MockFailTaskThatTimesOut <EOL> from ooni . tests . mocks import MockTimeoutOnceTask , MockFailTaskWithTimeout <EOL> from ooni . tests . mocks import MockTaskManager , mockFailure , MockDirector <EOL> from ooni . tests . mocks import MockNetTest , MockSuccessMeasurement <EOL> from ooni . tests . mocks import MockFailMeasurement <EOL> from ooni . settings import config <EOL> class TestTaskManager ( unittest . TestCase ) : <EOL> timeout = <NUM_LIT:1> <EOL> def setUp ( self ) : <EOL> self . measurementManager = MockTaskManager ( ) <EOL> self . measurementManager . concurrency = <NUM_LIT:20> <EOL> self . measurementManager . retries = <NUM_LIT:2> <EOL> self . measurementManager . start ( ) <EOL> self . clock = task . Clock ( ) <EOL> data_dir = os . path . dirname ( os . path . abspath ( __file__ ) ) <EOL> data_dir = os . path . join ( data_dir , '<STR_LIT:..>' , '<STR_LIT:..>' , '<STR_LIT:data>' ) <EOL> self . old_datadir = "<STR_LIT>" <EOL> if hasattr ( config . global_options , '<STR_LIT>' ) : <EOL> self . old_datadir = config . global_options [ '<STR_LIT>' ] <EOL> config . global_options [ '<STR_LIT>' ] = data_dir <EOL> config . set_paths ( ) <EOL> def tearDown ( self ) : <EOL> if self . old_datadir == "<STR_LIT>" : <EOL> del config . global_options [ '<STR_LIT>' ] <EOL> else : <EOL> config . global_options [ '<STR_LIT>' ] = self . old_datadir <EOL> config . set_paths ( ) <EOL> def schedule_successful_tasks ( self , task_type , number = <NUM_LIT:1> ) : <EOL> all_done = [ ] <EOL> for x in range ( number ) : <EOL> mock_task = task_type ( ) <EOL> all_done . append ( mock_task . done ) <EOL> self . measurementManager . schedule ( mock_task ) <EOL> d = defer . DeferredList ( all_done ) <EOL> @ d . addCallback <EOL> def done ( res ) : <EOL> for task_result , task_instance in self . measurementManager . successes : <EOL> self . assertEqual ( task_result , <NUM_LIT> ) <EOL> self . assertIsInstance ( task_instance , task_type ) <EOL> return d <EOL> def schedule_failing_tasks ( self , task_type , number = <NUM_LIT:1> ) : <EOL> all_done = [ ] <EOL> for x in range ( number ) : <EOL> mock_task = task_type ( ) <EOL> all_done . append ( mock_task . done ) <EOL> mock_task . done . addErrback ( lambda x : None ) <EOL> self . measurementManager . schedule ( mock_task ) <EOL> d = defer . DeferredList ( all_done ) <EOL> @ d . addCallback <EOL> def done ( res ) : <EOL> self . assertEqual ( self . measurementManager . failures , number * <NUM_LIT:3> ) <EOL> return d <EOL> def test_schedule_failing_with_mock_failure_task ( self ) : <EOL> mock_task = MockFailTask ( ) <EOL> self . measurementManager . schedule ( mock_task ) <EOL> self . assertFailure ( mock_task . done , MockFailure ) <EOL> return mock_task . done <EOL> def test_schedule_successful_one_task ( self ) : <EOL> return self . schedule_successful_tasks ( MockSuccessTask ) <EOL> def test_schedule_successful_one_task_with_timeout ( self ) : <EOL> return self . schedule_successful_tasks ( MockSuccessTaskWithTimeout ) <EOL> def test_schedule_failing_tasks_that_timesout ( self ) : <EOL> self . measurementManager . retries = <NUM_LIT:0> <EOL> task_type = MockFailTaskThatTimesOut <EOL> task_timeout = <NUM_LIT:5> <EOL> mock_task = task_type ( ) <EOL> mock_task . timeout = task_timeout <EOL> mock_task . clock = self . clock <EOL> self . measurementManager . schedule ( mock_task ) <EOL> self . clock . advance ( task_timeout ) <EOL> @ mock_task . done . addBoth <EOL> def done ( res ) : <EOL> self . assertEqual ( self . measurementManager . failures , <NUM_LIT:1> ) <EOL> return mock_task . done <EOL> def test_schedule_time_out_once ( self ) : <EOL> task_type = MockTimeoutOnceTask <EOL> task_timeout = <NUM_LIT:5> <EOL> mock_task = task_type ( ) <EOL> mock_task . timeout = task_timeout <EOL> mock_task . clock = self . clock <EOL> self . measurementManager . schedule ( mock_task ) <EOL> self . clock . advance ( task_timeout ) <EOL> @ mock_task . done . addBoth <EOL> def done ( res ) : <EOL> self . assertEqual ( self . measurementManager . failures , <NUM_LIT:1> ) <EOL> for task_result , task_instance in self . measurementManager . successes : <EOL> self . assertEqual ( task_result , <NUM_LIT> ) <EOL> self . assertIsInstance ( task_instance , task_type ) <EOL> return mock_task . done <EOL> def test_schedule_failing_one_task ( self ) : <EOL> return self . schedule_failing_tasks ( MockFailTask ) <EOL> def test_schedule_failing_one_task_with_timeout ( self ) : <EOL> return self . schedule_failing_tasks ( MockFailTaskWithTimeout ) <EOL> def test_schedule_successful_ten_tasks ( self ) : <EOL> return self . schedule_successful_tasks ( MockSuccessTask , number = <NUM_LIT:10> ) <EOL> def test_schedule_failing_ten_tasks ( self ) : <EOL> return self . schedule_failing_tasks ( MockFailTask , number = <NUM_LIT:10> ) <EOL> def test_schedule_successful_27_tasks ( self ) : <EOL> return self . schedule_successful_tasks ( MockSuccessTask , number = <NUM_LIT> ) <EOL> def test_schedule_failing_27_tasks ( self ) : <EOL> return self . schedule_failing_tasks ( MockFailTask , number = <NUM_LIT> ) <EOL> def test_task_retry_and_succeed ( self ) : <EOL> mock_task = MockFailOnceTask ( ) <EOL> self . measurementManager . schedule ( mock_task ) <EOL> @ mock_task . done . addCallback <EOL> def done ( res ) : <EOL> self . assertEqual ( self . measurementManager . failures , <NUM_LIT:1> ) <EOL> self . assertEqual ( self . measurementManager . successes , <EOL> [ ( <NUM_LIT> , mock_task ) ] ) <EOL> return mock_task . done <EOL> def test_task_retry_and_succeed_56_tasks ( self ) : <EOL> """<STR_LIT>""" <EOL> all_done = [ ] <EOL> number = <NUM_LIT> <EOL> for x in range ( number ) : <EOL> mock_task = MockFailOnceTask ( ) <EOL> all_done . append ( mock_task . done ) <EOL> self . measurementManager . schedule ( mock_task ) <EOL> d = defer . DeferredList ( all_done ) <EOL> @ d . addCallback <EOL> def done ( res ) : <EOL> self . assertEqual ( self . measurementManager . failures , number ) <EOL> for task_result , task_instance in self . measurementManager . successes : <EOL> self . assertEqual ( task_result , <NUM_LIT> ) <EOL> self . assertIsInstance ( task_instance , MockFailOnceTask ) <EOL> return d <EOL> class TestMeasurementManager ( unittest . TestCase ) : <EOL> def setUp ( self ) : <EOL> mock_director = MockDirector ( ) <EOL> self . measurementManager = MeasurementManager ( ) <EOL> self . measurementManager . director = mock_director <EOL> self . measurementManager . concurrency = <NUM_LIT:10> <EOL> self . measurementManager . retries = <NUM_LIT:2> <EOL> self . measurementManager . start ( ) <EOL> self . mockNetTest = MockNetTest ( ) <EOL> def test_schedule_and_net_test_notified ( self , number = <NUM_LIT:1> ) : <EOL> mock_task = MockSuccessMeasurement ( self . mockNetTest ) <EOL> self . measurementManager . schedule ( mock_task ) <EOL> @ mock_task . done . addCallback <EOL> def done ( res ) : <EOL> self . assertEqual ( self . mockNetTest . successes , <EOL> [ <NUM_LIT> ] ) <EOL> self . assertEqual ( len ( self . mockNetTest . successes ) , <NUM_LIT:1> ) <EOL> return mock_task . done <EOL> def test_schedule_failing_one_measurement ( self ) : <EOL> mock_task = MockFailMeasurement ( self . mockNetTest ) <EOL> self . measurementManager . schedule ( mock_task ) <EOL> @ mock_task . done . addErrback <EOL> def done ( failure ) : <EOL> self . assertEqual ( self . measurementManager . failures , <NUM_LIT:3> ) <EOL> self . assertEqual ( failure , mockFailure ) <EOL> self . assertEqual ( len ( self . mockNetTest . successes ) , <NUM_LIT:0> ) <EOL> return mock_task . done </s>
