def __init__(self, mission_id=None, script=None, dbobj=None):
        self.mission = None
        self.id = None
        self.name = None
        self.blocks = None
        self.version = None
        self.schedule = None
        self.active = False
        self.run_count = 0
        self.created = 0
        self.event_types = self._get_event_types()

        if dbobj is not None:
            self._load(dbobj)

        elif script is not None:
            with session_scope() as session:
                dbobj = self._create(session, script)
                self._load(dbobj)

        elif mission_id is not None:
            with session_scope() as session:
                dbobj = session.query(MissionData).filter(MissionData.id == mission_id)
                if dbobj:
                    self._load(dbobj)

        # RUN PARAMETERS
        self.running = False
        self.job = None
        self.state = None
        self.current_step = None

        self.vars = {}
        # self.executor = DummyExecutor()
        host = app.config['IA_HOST']
        port = app.config['IA_PORT']
        self.executor = RestExecutor(self.name, host, port, timeout=self.DEFAULT_TIMEOUT)
        app.scheduler.add_listener(self._job_event_listener, EVENT_JOB_EXECUTED | EVENT_JOB_ERROR)
        app.jms_reader.add_listener(self.jms_listener)

        if self.active:
            self.state = Tags.MISSION
            self._schedule_mission()
 def setUp(self):
     self.executor = RestExecutor('test', 'test', 12345)
     self.response = {'cmd': 'cmd', 'type': 'type', 'value': 'value', 'time': time.time()}
     self.response_json = json.dumps(self.response)
class Mission(object):
    DEFAULT_TIMEOUT = 30000

    def __init__(self, mission_id=None, script=None, dbobj=None):
        self.mission = None
        self.id = None
        self.name = None
        self.blocks = None
        self.version = None
        self.schedule = None
        self.active = False
        self.run_count = 0
        self.created = 0
        self.event_types = self._get_event_types()

        if dbobj is not None:
            self._load(dbobj)

        elif script is not None:
            with session_scope() as session:
                dbobj = self._create(session, script)
                self._load(dbobj)

        elif mission_id is not None:
            with session_scope() as session:
                dbobj = session.query(MissionData).filter(MissionData.id == mission_id)
                if dbobj:
                    self._load(dbobj)

        # RUN PARAMETERS
        self.running = False
        self.job = None
        self.state = None
        self.current_step = None

        self.vars = {}
        # self.executor = DummyExecutor()
        host = app.config['IA_HOST']
        port = app.config['IA_PORT']
        self.executor = RestExecutor(self.name, host, port, timeout=self.DEFAULT_TIMEOUT)
        app.scheduler.add_listener(self._job_event_listener, EVENT_JOB_EXECUTED | EVENT_JOB_ERROR)
        app.jms_reader.add_listener(self.jms_listener)

        if self.active:
            self.state = Tags.MISSION
            self._schedule_mission()

    def _get_event_types(self):
        session = app.Session()
        d = {}
        for event_type in session.query(EventType).all():
            d[event_type.name] = event_type.id
        return d

    def _create(self, session, data):
        log.info('Creating mission in database')
        mission_dict = yaml.load(data)
        validate(mission_dict, mission_schema.Mission.get_schema())
        name = mission_dict['name']
        version = mission_dict['version']

        # check if mission already exists, create if not
        mission = session.query(MissionData).filter(MissionData.name == name).one_or_none()
        if mission is None:
            mission = MissionData(name=name)
            session.add(mission)

        script = session.query(Script).filter(Script.name == name).filter(Script.version == version).one_or_none()
        if script is None:
            script = Script(script=data, name=name, version=version, mission=mission)
            session.add(script)

        if script.script != data:
            log.info('%r %r', script.script, data)
            raise DuplicateScriptException

        session.commit()

        # Set the current script to this script
        mission.script = script
        session.commit()
        return mission

    def _load(self, dbobj):
        log.info('Loading mission from database')
        self.mission_txt = dbobj.script.script
        self.mission = yaml.load(self.mission_txt)
        self.id = dbobj.id
        self.name = dbobj.name
        self.active = dbobj.active
        validate(self.mission, mission_schema.Mission.get_schema())
        self.blocks = self._load_blocks()
        self.version = dbobj.script.version
        self.run_count = len(dbobj.runs)
        self.schedule = self.mission.get('schedule')
        self.created = dbobj.script.create_time
        self.description = self.mission.get('desc')

    def delete(self):
        with session_scope() as session:
            dbobj = self._get_dbobj(session)
            dbobj.script = None
            session.commit()

    def _get_dbobj(self, session):
        return session.query(MissionData).filter(MissionData.id == self.id).one()

    def __repr__(self):
        return repr(self.mission)

    def _load_blocks(self):
        d = {}
        for block in self.mission.get('blocks', []):
            d[block['label']] = block
        return d

    def _job_event_listener(self, event):
        if event.job_id.startswith(self.name):
            if Tags.SCHEDULE not in self.mission:
                with session_scope() as session:
                    dbobj = self._get_dbobj(session)
                    dbobj.active = False
                    session.commit()
                    self.active = False

    def jms_listener(self, source, event):
        """
        Listen for JMS events
        Trigger an execute if a target event is received
        :param event:
        :return:
        """

        schedule = self.mission.get(Tags.SCHEDULE, {})

        # check if the incoming event is a trigger to execute the mission
        if schedule.get('source') == source and schedule.get('event') == event:

            log.debug('Scheduling Mission to Run immediately...')

            # Schedule the mission to run immediately
            self._add_job()

    def _get_events(self, session, run_id=None):
        events = []
        dbobj = self._get_dbobj(session)
        if not dbobj.runs:
            return []

        if run_id is None:
            run = dbobj.runs[0]
        else:
            run = session.query(Run).filter(Run.mission_id == self.id).filter(Run.id == run_id).one_or_none()

        if run and run.events:
            for event in run.events[:10]:
                try:
                    e = json.loads(event.event)
                except ValueError:
                    e = event.event
                events.append((event.timestamp.isoformat(), event.type.name, e))
            return events
        return []

    def small(self):
        job = app.scheduler.get_job(self.name)
        next_run = job.next_run_time.isoformat() if job else None
        return {
            'id': self.id,
            'name': self.name,
            'version': self.version,
            'desc': self.description,
            'active': self.active,
            'running': self.running,
            'current_step': self.current_step,
            'run_count': self.run_count,
            'schedule': self.schedule,
            'next_run': next_run,
            'created': self.created.isoformat(),
        }

    def full(self):
        with session_scope() as session:
            events = self._get_events(session)
            base = self.small()
            base['events'] = events
            base['script'] = self.mission_txt
            return base

    def activate(self):
        if not self.active:
            with session_scope() as session:
                log.debug('Activating mission: %s', self.name)
                dbobj = self._get_dbobj(session)
                dbobj.active = True
                session.commit()

                self.active = True
                self.state = Tags.MISSION
                self._schedule_mission()

    def _add_job(self, trigger=None, kwargs=None):
        trigger = trigger or 'date'
        kwargs = kwargs or {}
        job_id = self.name
        app.scheduler.add_job(self._execute_mission, trigger, id=job_id, **kwargs)

    def _schedule_mission(self):
        """
        schedule the mission
        :return:
        """
        schedule = self.mission.get(Tags.SCHEDULE, {})
        cron_keys = {'year', 'month', 'day', 'week', 'day_of_week',
                     'hour', 'minute', 'second', 'start_date', 'end_date'}

        event_keys = {'source', 'event'}

        trigger = None

        if len(cron_keys.intersection(schedule.keys())) > 0:
            trigger = Tags.CRON
        elif len(event_keys.intersection(schedule.keys())) == 0:
            trigger = 'date'

        if trigger:
            self._add_job(trigger, schedule)

    def deactivate(self):
        if self.active:
            with session_scope() as session:
                log.debug('Deactivating mission: %s', self.name)
                dbobj = self._get_dbobj(session)
                dbobj.active = False
                job_id = self.name
                app.scheduler.remove_job(job_id)
                self.active = False

    def _add_event(self, session, run, event_type, event=''):
        if event_type not in self.event_types:
            et = EventType(name=event_type)
            session.add(et)
            session.commit()
            self.event_types[event_type] = et.id

        if not isinstance(event, basestring):
            try:
                event = json.dumps(event, cls=MyEncoder)
            except TypeError:
                log.error('Unable to create JSON from: %r %r', type(event), event)
                event = str(event)

        event = Event(run=run, event_type_id=self.event_types[event_type], event=event)
        session.add(event)
        session.commit()

    def _execute_mission(self):
        with session_scope() as session:
            dbobj = self._get_dbobj(session)
            run = Run(mission=dbobj, script=dbobj.script)
            session.add(run)
            session.commit()
            add_event = functools.partial(self._add_event, session, run)

            add_event('start')

            error_policy = ErrorPolicy(self.mission.get('onerror', {}))
            sequence = self.blocks.get(Tags.MISSION)
            if sequence is not None:
                complete = False
                attempt = 0
                max_attempts = error_policy.count
                backoff = error_policy.backoff
                while not complete and attempt < max_attempts:
                    try:
                        with lock_instrument(self.mission.get(Tags.INSTRUMENT, []), self.executor, add_event):
                            self.running = True
                            self.run_count += 1
                            self._execute_sequence(Tags.MISSION, add_event)
                            complete = True
                    except (LockException, ConnectionError) as e:
                        log.error('Exception locking instruments for mission: %s (%r)', self.name, e)
                        if error_policy.action == 'retry':
                            attempt += 1
                            time.sleep(backoff)
                        else:
                            break
                    except (InstrumentException, PolicyException, CommandArgumentException, ConnectionError) as e:
                        log.error('Exception when processing mission, aborting mission (%r)', e)
                        add_event('exception', str(e))
                        break
                    finally:
                        self.running = False
                        self.current_step = None

                if not complete:
                    log.error('Unable to complete mission: %s', self.name)

            add_event('completion')

    def _execute_sequence(self, section, add_event):
        block = self.blocks.get(section)
        sequence = block.get('sequence', [])
        block_error_policy = ErrorPolicy(block.get('onerror', self.mission.get('onerror', {})))
        if sequence is not None:
            for index, step in enumerate(sequence):
                error_policy = ErrorPolicy(step.get('onerror', {})) if 'onerror' in block else block_error_policy
                self.current_step = (index, step)
                log.info('Executing step: %s from mission: %s section: %s', step, self.name, section)
                add_event('step', step)
                rval = self._handle_step(step, error_policy, add_event)
                if rval is not None:
                    add_event('result', rval)

    def _handle_step(self, step, error_policy, add_event):
        log.info('step: %r', step)
        count = 0
        while count < error_policy.count:
            count += 1
            try:
                if 'block_name' in step:
                    if self._eval_conditional(step):
                        loop = step.get('loop', 1)
                        for _ in xrange(loop):
                            self._execute_sequence(step['block_name'], add_event)
                    return

                if 'sleep' in step:
                    return time.sleep(step['sleep'])

                rval = self.executor.command(step)
                if step.get('get_state'):
                    self.vars['driver_state'] = rval.value
                elif step.get('get') and step.get('parameter'):
                    self.vars[step.get('parameter')] = rval.value
                return rval

            except Exception as e:
                if error_policy.action == 'abort':
                    raise e
                if error_policy.action == 'continue':
                    log.error('Exception in step: %r executing continue policy: %r', step, e)
                    return

        raise PolicyException

    def _eval_conditional(self, step):
        log.debug('eval_conditional: %r', step)
        conditional = step.get('condition')
        if not conditional:
            return True

        expected_value = conditional.get('value')
        current_value = self.vars.get(conditional.get('variable'))
        comparator = conditional.get('comparator', 'equal')
        log.debug('eval_conditional %r %r %r', expected_value, comparator, current_value)
        return (comparator == 'equal') ^ (not expected_value == current_value)

    @staticmethod
    def from_script(data):
        return Mission(script=data)

    @staticmethod
    def load_all():
        log.info('LOADING ALL MISSIONS FROM DATABASE')
        session = app.Session()
        missions = session.query(MissionData).all()
        d = {}
        for m in missions:
            if m.script is not None:
                m = Mission(dbobj=m)
                d[m.id] = m
        return d

    def versions(self):
        with session_scope() as session:
            scripts = session.query(Script).filter(Script.mission_id == self.id).all()
            return [script.id for script in scripts]

    def get_version(self, version_id):
        with session_scope() as session:
            script = session.query(Script).filter(Script.mission_id == self.id)\
                .filter(Script.id == version_id).one_or_none()
            if script is not None:
                return script.script

    def set_version(self, version_id):
        with session_scope() as session:
            script = session.query(Script).filter(Script.mission_id == self.id)\
                .filter(Script.id == version_id).one_or_none()
            if script is None:
                return False

            dbobj = self._get_dbobj(session)
            dbobj.script = script
            self.version = dbobj.script.version
            session.commit()
            return True

    def runs(self):
        with session_scope() as session:
            runs = session.query(Run).filter(Run.mission_id == self.id).all()
            runs = [run.id for run in runs]
            return runs

    def get_run(self, run_id):
        with session_scope() as session:
            return self._get_events(session, run_id)
class RestExecutorUnitTest(unittest.TestCase):
    def setUp(self):
        self.executor = RestExecutor('test', 'test', 12345)
        self.response = {'cmd': 'cmd', 'type': 'type', 'value': 'value', 'time': time.time()}
        self.response_json = json.dumps(self.response)

    def assert_response(self, response):
        self.assertEqual(response.cmd, self.response.get('cmd'))
        self.assertEqual(response.type, self.response.get('type'))
        self.assertEqual(response.value, self.response.get('value'))
        self.assertEqual(response.time, self.response.get('time'))
        self.assertEqual(response.status_code, 200)

    def test_url(self):
        expected = 'http://test:12345/instrument/api/target/name'
        returned = self.executor._url('target', 'name')
        self.assertEqual(expected, returned)

    @httpretty.activate
    def test_execute_resource(self):
        httpretty.register_uri(httpretty.POST, self.executor._url('target', 'execute'), body=self.response_json)
        response = self.executor.execute_resource('target', 'test', {'key1': 'value1'}, timeout=60000)
        self.assert_response(response)

    @httpretty.activate
    def test_reset(self):
        httpretty.register_uri(httpretty.POST, self.executor._url('target', 'shutdown'), body=self.response_json)
        response = self.executor.reset('target', timeout=60000)
        self.assert_response(response)

    @httpretty.activate
    def test_ping(self):
        httpretty.register_uri(httpretty.POST, self.executor._url('target', 'ping'), body=self.response_json)
        response = self.executor.ping('target', timeout=60000)
        self.assert_response(response)

    @httpretty.activate
    def test_discover(self):
        httpretty.register_uri(httpretty.POST, self.executor._url('target', 'discover'), body=self.response_json)
        response = self.executor.discover('target', timeout=60000)
        self.assert_response(response)

    @httpretty.activate
    def test_get_state(self):
        httpretty.register_uri(httpretty.GET, self.executor._url('target', 'state'), body=self.response_json)
        response = self.executor.get_state('target', timeout=60000)
        self.assert_response(response)

    @httpretty.activate
    def test_get_resource(self):
        httpretty.register_uri(httpretty.GET, self.executor._url('target', 'resource'), body=self.response_json)
        response = self.executor.get_resource('target', 'parameter', timeout=60000)
        self.assert_response(response)

    @httpretty.activate
    def test_set_resource(self):
        httpretty.register_uri(httpretty.POST, self.executor._url('target', 'resource'), body=self.response_json)
        response = self.executor.set_resource('target', {}, timeout=60000)
        self.assert_response(response)

    @httpretty.activate
    def test_disconnect(self):
        httpretty.register_uri(httpretty.POST, self.executor._url('target', 'disconnect'), body=self.response_json)
        response = self.executor.disconnect('target', timeout=60000)
        self.assert_response(response)

    @httpretty.activate
    def test_connect(self):
        httpretty.register_uri(httpretty.POST, self.executor._url('target', 'connect'), body=self.response_json)
        response = self.executor.connect('target', timeout=60000)
        self.assert_response(response)

    @httpretty.activate
    def test_set_init_params(self):
        httpretty.register_uri(httpretty.POST, self.executor._url('target', 'initparams'), body=self.response_json)
        response = self.executor.set_init_params('target', {}, timeout=60000)
        self.assert_response(response)

    @httpretty.activate
    def test_lock(self):
        httpretty.register_uri(httpretty.POST, self.executor._url('target', 'lock'), body=self.response_json)
        self.executor.lock(['target'])

    @httpretty.activate
    def test_unlock(self):
        httpretty.register_uri(httpretty.GET, self.executor._url('target', 'lock'), body='"test"')
        httpretty.register_uri(httpretty.POST, self.executor._url('target', 'unlock'), body=self.response_json)
        self.executor.unlock(['target'])