def test_failed_job_lifecycle(self):
        with patch('openquake.job.Job.from_file') as from_file:

            # called in place of Job.launch
            def test_status_running_and_fail():
                self.assertEquals('running', self._job_status())

                raise Exception('OMG!')

            # replaces Job.launch with a mock
            def patch_job_launch(*args, **kwargs):
                self.job = self.job_from_file(*args, **kwargs)
                self.job.launch = mock.Mock(
                    side_effect=test_status_running_and_fail)

                self.assertEquals('pending', self._job_status())

                return self.job

            from_file.side_effect = patch_job_launch

            with patch('openquake.job.spawn_job_supervisor'):
                self.assertRaises(Exception, run_job,
                                helpers.get_data_path(CONFIG_FILE), 'db')

        self.assertEquals(1, self.job.launch.call_count)
        self.assertEquals('failed', self._job_status())

    def test_failed_job_lifecycle(self):
        def test_status_running_and_fail(*args):
            self.assertEquals("running", self._calculation_status())

            raise Exception("OMG!")

        def patch_job_launch(*args, **kwargs):
            self.job = self.job_from_file(*args, **kwargs)

            self.assertEquals("pending", self._calculation_status())

            return self.job

        before_launch = engine._launch_calculation
        try:
            engine._launch_calculation = mock.Mock(side_effect=test_status_running_and_fail)

            with patch("openquake.engine._job_from_file") as from_file:
                from_file.side_effect = patch_job_launch

                with patch("os.fork", mocksignature=False) as fork:
                    fork.return_value = 0
                    self.assertRaises(Exception, engine.run_calculation, self.calc_proxy, self.params, self.sections)

            self.assertEquals(1, engine._launch_calculation.call_count)
            self.assertEquals("failed", self._calculation_status())
        finally:
            engine._launch_calculation = before_launch

    def test_successful_job_lifecycle(self):
        with patch('openquake.job.Job.from_file') as from_file:

            # called in place of Job.launch
            def test_status_running_and_succeed():
                self.assertEquals('running', self._job_status())

                return []

            # replaces Job.launch with a mock
            def patch_job_launch(*args, **kwargs):
                self.job = self.job_from_file(*args, **kwargs)
                self.job.launch = mock.Mock(
                    side_effect=test_status_running_and_succeed)

                self.assertEquals('pending', self._job_status())

                return self.job

            from_file.side_effect = patch_job_launch

            with patch('openquake.job.spawn_job_supervisor'):
                run_job(helpers.get_data_path(CONFIG_FILE), 'db')

        self.assertEquals(1, self.job.launch.call_count)
        self.assertEquals('succeeded', self._job_status())

    def setUp(self):
        self.job = engine.prepare_job()
        self.calc = disagg_core.DisaggHazardCalculator(self.job)

        # Mock `disagg_task_arg_gen`
        disagg_path = 'openquake.engine.calculators.hazard.disaggregation'
        self.disagg_tag_patch = helpers.patch(
            '%s.core.DisaggHazardCalculator.disagg_task_arg_gen'
            % disagg_path)
        self.disagg_tag_mock = self.disagg_tag_patch.start()
        # fake disagg task arg generator:
        disagg_tag = iter(xrange(3))
        self.disagg_tag_mock.return_value = disagg_tag

        # Mock `haz_general.queue_next`
        base_path = 'openquake.engine.calculators.base'
        self.queue_next_patch = helpers.patch('%s.queue_next' % base_path)
        self.queue_next_mock = self.queue_next_patch.start()

        # Mock `finalize_hazard_curves`
        general_path = 'openquake.engine.calculators.hazard.general'
        self.finalize_curves_patch = helpers.patch(
            '%s.BaseHazardCalculator.finalize_hazard_curves'
            % general_path)
        self.finalize_curves_mock = self.finalize_curves_patch.start()

    def test_successful_job_lifecycle(self):
        def test_status_running_and_succeed(*args):
            self.assertEqual("running", self._calculation_status())

            return []

        def patch_job_launch(*args, **kwargs):
            self.job = self.job_from_file(*args, **kwargs)

            self.assertEqual("pending", self._calculation_status())

            return self.job

        before_launch = engine._launch_job
        try:
            engine._launch_job = mock.Mock(side_effect=test_status_running_and_succeed)

            with patch("openquake.engine._job_from_file") as from_file:
                from_file.side_effect = patch_job_launch

                with patch("os.fork", mocksignature=False) as fork:
                    fork.return_value = 0
                    engine.run_job(self.job, self.params, self.sections)

            self.assertEqual(1, engine._launch_job.call_count)
            self.assertEqual("succeeded", self._calculation_status())
        finally:
            engine._launch_job = before_launch

    def test_job_launch_calls_record_initial_stats(self):
        '''When a job is launched, make sure that
        :py:method:`openquake.engine.JobContext._record_initial_stats`
        is called.
        '''
        # Mock out pieces of the test job so it doesn't actually run.
        eb_haz_calc = ('openquake.calculators.hazard.event_based.core'
                       '.EventBasedHazardCalculator')
        eb_risk_calc = ('openquake.calculators.risk.event_based.core'
                       '.EventBasedRiskCalculator')
        methods = ('initialize', 'pre_execute', 'execute', 'post_execute')

        haz_patchers = [patch('%s.%s' % (eb_haz_calc, m)) for m in methods]
        risk_patchers = [patch('%s.%s' % (eb_risk_calc, m)) for m in methods]

        for p in haz_patchers:
            p.start()
        for p in risk_patchers:
            p.start()

        try:
            record = 'openquake.engine.JobContext._record_initial_stats'

            with patch(record) as record_mock:
                engine._launch_job(
                    self.eb_job, ['general', 'HAZARD', 'RISK'])

                self.assertEqual(1, record_mock.call_count)
        finally:
            for p in haz_patchers:
                p.stop()
            for p in risk_patchers:
                p.stop()

 def setUp(self):
     self.db_patch = patch(
         'openquake.engine.utils.monitor._db_cnode_status')
     self.live_patch = (
         patch('openquake.engine.utils.monitor._live_cnode_status'))
     self.db_mock = self.db_patch.start()
     self.live_mock = self.live_patch.start()

    def test_job_launch_calls_record_initial_stats(self):
        """When a job is launched, make sure that
        :py:method:`openquake.engine.JobContext._record_initial_stats`
        is called.
        """
        # Mock out pieces of the test job so it doesn't actually run.
        eb_haz_calc = "openquake.calculators.hazard.event_based.core" ".EventBasedHazardCalculator"
        eb_risk_calc = "openquake.calculators.risk.event_based.core" ".EventBasedRiskCalculator"
        methods = ("initialize", "pre_execute", "execute", "post_execute")

        haz_patchers = [patch("%s.%s" % (eb_haz_calc, m)) for m in methods]
        risk_patchers = [patch("%s.%s" % (eb_risk_calc, m)) for m in methods]

        for p in haz_patchers:
            p.start()
        for p in risk_patchers:
            p.start()

        try:
            record = "openquake.engine.JobContext._record_initial_stats"

            with patch(record) as record_mock:
                engine._launch_job(self.eb_job, ["general", "HAZARD", "RISK"])

                self.assertEqual(1, record_mock.call_count)
        finally:
            for p in haz_patchers:
                p.stop()
            for p in risk_patchers:
                p.stop()

    def test_pre_execute(self):
        # Most of the pre-execute functionality is implement in other methods.
        # For this test, just make sure each method gets called.
        path = ('openquake.engine.calculators.risk.general.BaseRiskCalculator')
        patches = (
            helpers.patch(
                '%s.%s' % (path, '_store_exposure')),
            helpers.patch(
                '%s.%s' % (path, 'set_risk_models')),
            helpers.patch(
                '%s.%s' % (path, '_initialize_progress')))

        mocks = [p.start() for p in patches]

        mocks[0].return_value = mock.Mock()
        mocks[0].return_value.taxonomies_in.return_value = {'RC': 10}

        self.calculator.imt = 'PGA'
        self.calculator.pre_execute()

        for i, m in enumerate(mocks):
            self.assertEqual(1, m.call_count,
                             "mock %d has not been called" % (i + 1))
            m.stop()
            patches[i].stop()

 def setUp(self):
     self.monitor_patch = patch(
         "openquake.utils.monitor.count_failed_nodes")
     self.stats_patch = patch(
         "openquake.utils.stats.get_progress_timing_data")
     self.monitor_mock = self.monitor_patch.start()
     self.stats_mock = self.stats_patch.start()

    def test_compute_uhs_task_pi(self):
        # Test that progress indicators are working properly for
        # `compute_uhs_task`.

        # Mock out the two 'heavy' functions called by this task;
        # we don't need to do these and we don't want to waste the cycles.
        cmpt_uhs = '%s.%s' % (self.UHS_CORE_MODULE, 'compute_uhs')
        write_uhs_data = '%s.%s' % (self.UHS_CORE_MODULE,
                                    'write_uhs_spectrum_data')
        with helpers.patch(cmpt_uhs):
            with helpers.patch(write_uhs_data):

                get_counter = lambda: stats.get_counter(
                    self.job_id, 'h', 'compute_uhs_task', 'i')

                # First, check that the counter for `compute_uhs_task` is
                # `None`:
                self.assertIsNone(get_counter())

                realization = 0
                site = Site(0.0, 0.0)
                # execute the task as a plain old function
                compute_uhs_task(self.job_id, realization, site)
                self.assertEqual(1, get_counter())

                compute_uhs_task(self.job_id, realization, site)
                self.assertEqual(2, get_counter())

    def test_cleanup_after_job(self):
        with patch('openquake.engine.kvs.cache_gc') as cache_gc:
            with patch('openquake.engine.supervising.supervisor.'
                       '_get_task_ids') as gti:
                with patch('celery.task.control.revoke') as revoke:
                    gti.return_value = ['task-id-1', 'task-id-2']

                    supervisor.cleanup_after_job(self.job.id, terminate=True)

                    self.assertEqual(1, cache_gc.call_count)
                    self.assertEqual(((self.job.id, ), {}), cache_gc.call_args)

                    self.assertEqual(1, gti.call_count)
                    self.assertEqual(((self.job.id, ), {}), gti.call_args)

                    self.assertEqual(2, revoke.call_count)
                    exp_revoke_args = [(('task-id-1',), {'terminate': True}),
                                       (('task-id-2',), {'terminate': True})]
                    self.assertEqual(exp_revoke_args, revoke.call_args_list)

                with patch('celery.task.control.revoke') as revoke:
                    gti.return_value = ['task-id-1', 'task-id-2']

                    supervisor.cleanup_after_job(self.job.id, terminate=False)

                    self.assertEqual(2, cache_gc.call_count)
                    self.assertEqual(((self.job.id, ), {}), cache_gc.call_args)

                    self.assertEqual(2, gti.call_count)
                    self.assertEqual(((self.job.id, ), {}), gti.call_args)

                    self.assertEqual(2, revoke.call_count)
                    exp_revoke_args = [(('task-id-1',), {'terminate': False}),
                                       (('task-id-2',), {'terminate': False})]
                    self.assertEqual(exp_revoke_args, revoke.call_args_list)

    def test_pre_execute(self):
        base_path = ('openquake.engine.calculators.hazard.disaggregation.core'
                     '.DisaggHazardCalculator')
        init_src_patch = helpers.patch(
            '%s.%s' % (base_path, 'initialize_sources'))
        init_rlz_patch = helpers.patch(
            '%s.%s' % (base_path, 'initialize_realizations'))
        record_stats_patch = helpers.patch(
            '%s.%s' % (base_path, 'record_init_stats'))
        init_pr_data_patch = helpers.patch(
            '%s.%s' % (base_path, 'initialize_pr_data'))
        patches = (init_src_patch, init_rlz_patch,
                   record_stats_patch, init_pr_data_patch)

        mocks = [p.start() for p in patches]

        self.calc.pre_execute()

        # make sure the site_collection is loaded:
        self.assertIsNotNone(self.calc.hc.site_collection)

        for i, m in enumerate(mocks):
            self.assertEqual(1, m.call_count)
            m.stop()
            patches[i].stop()

    def test_compute_uhs_with_site_model(self):
        the_job = helpers.prepare_job_context(
            helpers.demo_file('uhs/config_with_site_model.gem'))
        the_job.to_kvs()

        site = Site(0, 0)

        helpers.store_hazard_logic_trees(the_job)

        get_sm_patch = helpers.patch(
            'openquake.calculators.hazard.general.get_site_model')
        get_closest_patch = helpers.patch(
            'openquake.calculators.hazard.general.get_closest_site_model_data')
        compute_patch = helpers.patch(
            'openquake.calculators.hazard.uhs.core._compute_uhs')

        get_sm_mock = get_sm_patch.start()
        get_closest_mock = get_closest_patch.start()
        compute_mock = compute_patch.start()

        get_closest_mock.return_value = SiteModel(
            vs30=800, vs30_type='measured', z1pt0=100, z2pt5=200)
        try:
            compute_uhs(the_job, site)

            self.assertEqual(1, get_sm_mock.call_count)
            self.assertEqual(1, get_closest_mock.call_count)
            self.assertEqual(1, compute_mock.call_count)
        finally:
            get_sm_patch.stop()
            get_closest_patch.stop()
            compute_patch.stop()

    def test_pre_execute(self):
        # Most of the pre-execute functionality is implement in other methods.
        # For this test, just make sure each method gets called.
        base_path = ('openquake.engine.calculators.hazard.classical.core'
                     '.ClassicalHazardCalculator')
        init_src_patch = helpers.patch(
            '%s.%s' % (base_path, 'initialize_sources'))
        init_sm_patch = helpers.patch(
            '%s.%s' % (base_path, 'initialize_site_model'))
        init_rlz_patch = helpers.patch(
            '%s.%s' % (base_path, 'initialize_realizations'))
        record_stats_patch = helpers.patch(
            '%s.%s' % (base_path, 'record_init_stats'))
        init_pr_data_patch = helpers.patch(
            '%s.%s' % (base_path, 'initialize_pr_data'))
        patches = (init_src_patch, init_sm_patch, init_rlz_patch,
                   record_stats_patch, init_pr_data_patch)

        mocks = [p.start() for p in patches]

        # we don't expect the site collection to be loaded yet:
        self.assertIsNone(self.calc.hc._site_collection)

        self.calc.pre_execute()

        # make sure the site_collection is loaded:
        self.assertIsNotNone(self.calc.hc._site_collection)

        for i, m in enumerate(mocks):
            self.assertEqual(1, m.call_count)
            m.stop()
            patches[i].stop()

    def test_failed_job_lifecycle(self):

        def test_status_running_and_fail(*args):
            self.assertEqual('running', self._calculation_status())

            raise Exception('OMG!')

        def patch_job_launch(*args, **kwargs):
            self.job = self.job_from_file(*args, **kwargs)

            self.assertEqual('pending', self._calculation_status())

            return self.job

        before_launch = engine._launch_job
        try:
            engine._launch_job = mock.Mock(
                side_effect=test_status_running_and_fail)

            with patch('openquake.engine._job_from_file') as from_file:
                from_file.side_effect = patch_job_launch

                with patch('os.fork', mocksignature=False) as fork:
                    fork.return_value = 0
                    self.assertRaises(Exception, engine.run_job,
                                      self.job, self.params, self.sections)

            self.assertEqual(1, engine._launch_job.call_count)
            self.assertEqual('failed', self._calculation_status())
        finally:
            engine._launch_job = before_launch

    def test_loss_map_not_serialized_unless_conditional_loss_poes(self):
        with patch('openquake.risk.job.probabilistic'
                   '.aggregate_loss_curve.plot_aggregate_curve'):
            with patch('openquake.output.risk.create_loss_map_writer') as clw:
                clw.return_value = None

                self.mixin.execute()
                self.assertFalse(clw.called)

    def test_loss_map_serialized_if_conditional_loss_poes(self):
        self.mixin.params['CONDITIONAL_LOSS_POE'] = '0.01 0.02'

        with patch('openquake.risk.job.probabilistic'
                   '.aggregate_loss_curve.plot_aggregate_curve'):
            with patch('openquake.output.risk.create_loss_map_writer') as clw:
                clw.return_value = None

                self.mixin.execute()
                self.assertTrue(clw.called)

 def test_jvm_memmax_setting_is_not_passed(self):
     """Do not pass -Xmx to the jvm."""
     with helpers.patch("jpype.startJVM") as startjvm_mock:
         with helpers.patch("jpype.isJVMStarted") as isjvmstarted_mock:
             # Make sure that startJVM() gets called.
             isjvmstarted_mock.side_effect = lambda: False
             with helpers.patch("openquake.java.init_logs"):
                 java.jvm()
                 args, _ = startjvm_mock.call_args
                 self.assertFalse(filter(lambda a: a.startswith("-Xmx"), args))

    def test_loss_map_not_serialized_unless_conditional_loss_poes(self):
        with helpers.patch('openquake.calculators.risk.event_based.core'
                           '.plot_aggregate_curve'):
            with helpers.patch(
                'openquake.output.risk.create_loss_map_writer') as clw:

                clw.return_value = None

                self.calculator.execute()
                self.assertFalse(clw.called)