<commit_before>"""ParametricJobs Table."""
import os
import re
import time
import json
import logging
import calendar
from datetime import datetime

import cherrypy
from sqlalchemy import Column, Integer, Boolean, String, PickleType, TIMESTAMP, ForeignKey, Enum
from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy.orm import relationship
from sqlalchemy.orm.exc import NoResultFound, MultipleResultsFound

from rpc.DiracRPCClient import dirac_api_client, ParametricDiracJobClient
from utils.collections import list_splitter
from utils.tempfile_utils import temporary_runscript, temporary_macro
from ..utils import db_session
from ..statuses import LOCALSTATUS
from .SQLTableBase import SQLTableBase
from .JSONTableEncoder import JSONTableEncoder
from .DiracJobs import DiracJobs

logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
UNIXDATE = re.compile(r'(?P<month>[0-9]{2})-(?P<day>[0-9]{2})-(?P<year>[0-9]{4})$')

@cherrypy.expose
class ParametricJobs(SQLTableBase):
    """Jobs SQL Table."""

    __tablename__ = 'parametricjobs'
    id = Column(Integer, primary_key=True)  # pylint: disable=invalid-name
    app = Column(String(250))
    app_version = Column(String(250))
    sim_lfn_outputdir = Column(String(250))
    mctruth_lfn_outputdir = Column(String(250))
    macro = Column(String(250))
    tag = Column(String(250))
    njobs = Column(Integer)
    nevents = Column(Integer)
    seed = Column(Integer)
    fastnest_version = Column(String(250))
    reduction_version = Column(String(250))
    reduction_lfn_inputdir = Column(String(250))
    reduction_lfn_outputdir = Column(String(250))
    der_version = Column(String(250))
    der_lfn_inputdir = Column(String(250))
    der_lfn_outputdir = Column(String(250))
    lzap_version = Column(String(250))
    lzap_lfn_inputdir = Column(String(250))
    lzap_lfn_outputdir = Column(String(250))
    request_id = Column(Integer, ForeignKey('requests.id'), nullable=False)
    request = relationship("Requests", back_populates="parametricjobs")
    status = Column(Enum(LOCALSTATUS), nullable=False)
    reschedule = Column(Boolean, nullable=False, default=False)
    timestamp = Column(TIMESTAMP, nullable=False, default=datetime.utcnow, onupdate=datetime.utcnow)
    num_completed = Column(Integer, nullable=False, default=0)
    num_failed = Column(Integer, nullable=False, default=0)
    num_submitted = Column(Integer, nullable=False, default=0)
    num_running = Column(Integer, nullable=False, default=0)
    diracjobs = relationship("DiracJobs", back_populates="parametricjob")

    @hybrid_property
    def num_other(self):
        return self.njobs - (self.num_submitted + self.num_running + self.num_failed + self.num_completed)

    def submit(self):
        """Submit parametric job."""
        # problems here if not running simulation, there will be no macro so
        # everything including context needs reworking.
#        lfn_root = os.path.join('/lz/user/l/lzproduser.grid.hep.ph.ic.ac.uk', '_'.join(('-'.join((self.app, self.app_version)),
#                                                           '-'.join(('DER', self.der_version)))))

        parametric_job = ParametricDiracJobClient()
        dirac_ids = set()
        if self.app_version is not None:

            macro_name = os.path.splitext(os.path.basename(self.macro or ''))[0]
            livetime_sec_per_beamon = 0.1132698957
            livetimeperjob = str((self.nevents or 1) * livetime_sec_per_beamon)
            unixtime = time.time()
            match = UNIXDATE.search(macro_name)
            if match is not None:
                month, day, year = match.groups()
                unixtime = str(int(calendar.timegm(datetime(int(year), int(month), int(day), 0, 0).utctimetuple())))
            with temporary_runscript(root_version='5.34.32',
                                     root_arch='slc6_gcc44_x86_64',
                                     g4_version='4.9.5.p02',
                                     physics_version='1.4.0',
                                     se='UKI-LT2-IC-HEP-disk',
                                     unixtime=unixtime,
                                     livetimeperjob=livetimeperjob, **self) as runscript,\
                 temporary_macro(self.tag, self.macro or '', self.app, self.nevents) as macro:
                logger.info("Submitting ParametricJob %s, macro: %s to DIRAC", self.id, self.macro)
                for sublist in list_splitter(range(self.seed, self.seed + self.njobs), 1000):
                    with parametric_job as j:
                        j.setName(os.path.splitext(os.path.basename(macro))[0] + '-%(args)s')
                        j.setPlatform('ANY')
                        j.setExecutable(os.path.basename(runscript),
                                        os.path.basename(macro) + ' %(args)s',
                                        'lzproduction_output.log')
                        j.setInputSandbox([runscript, macro])
                        j.setDestination('ANY')
                        j.setParameterSequence('args', sublist, addToWorkflow=False)
                    dirac_ids.update(parametric_job.subjob_ids)
        else:
            with temporary_runscript(root_version='5.34.32',
                                     root_arch='slc6_gcc44_x86_64',
                                     g4_version='4.9.5.p02',
                                     physics_version='1.4.0',
                                     se='UKI-LT2-IC-HEP-disk', **self) as runscript:
                logger.info("Submitting ParametricJob %s, inputdir: %s to DIRAC", self.id, self.reduction_lfn_inputdir or self.der_lfn_inputdir or self.lzap_lfn_inputdir)

                input_lfn_dir=self.reduction_lfn_inputdir or\
                               self.der_lfn_inputdir or \
                               self.lzap_lfn_inputdir
                for sublist in list_splitter(list_lfns(input_lfn_dir), 1000):
                    with parametric_job as j:
                        j.setName("%(args)s")
                        j.setPlatform('ANY')
                        j.setExecutable(os.path.basename(runscript),
                                        '%(args)s',
                                        'lzanalysis_output.log')
                        j.setInputSandbox([runscript])
                        j.setDestination('ANY')
                        j.setParameterSequence('InputData', sublist, addToWorkflow='ParametricInputData')
                        j.setParameterSequence('args',
                                               [os.path.basename(l) for l in sublist],
                                               addToWorkflow=False)
                    dirac_ids.update(parametric_job.subjob_ids)

        with db_session() as session:
            session.bulk_insert_mappings(DiracJobs, [{'id': i, 'parametricjob_id': self.id}
                                                     for i in parametric_job.subjob_ids])

    def reset(self):
        """Reset parametric job."""
        with db_session(reraise=False) as session:
            dirac_jobs = session.query(DiracJobs).filter_by(parametricjob_id=self.id)
            dirac_job_ids = [j.id for j in dirac_jobs.all()]
            dirac_jobs.delete(synchronize_session=False)
        with dirac_api_client() as dirac:
            logger.info("Removing Dirac jobs %s from ParametricJob %s", dirac_job_ids, self.id)
            dirac.kill(dirac_job_ids)
            dirac.delete(dirac_job_ids)

    def delete_dirac_jobs(self, session):
        """Delete associated DIRAC jobs."""
        logger.info("Deleting DiracJobs for ParametricJob id: %s", self.id)
        session.query(DiracJobs)\
               .filter_by(parametricjob_id=self.id)\
               .delete(synchronize_session=False)

    def update_status(self):
        """Update the status of parametric job."""
        local_statuses = DiracJobs.update_status(self)
        # could just have DiracJobs return this... maybe better
#        local_statuses = Counter(status.local_status for status in dirac_statuses.elements())
        status = max(local_statuses or [self.status])
        with db_session() as session:
            this = session.merge(self)
            this.status = status
            this.num_completed = local_statuses['Completed']
            this.num_failed = local_statuses['Failed']
            this.num_submitted = local_statuses['Submitted']
            this.num_running = local_statuses['Running']
            this.reschedule = False
        return status

    @staticmethod
    def GET(reqid):  # pylint: disable=invalid-name
        """
        REST Get method.

        Returns all ParametricJobs for a given request id.
        """
        logger.debug("In GET: reqid = %s", reqid)
        requester = cherrypy.request.verified_user
        with db_session() as session:
            user_requests = session.query(ParametricJobs)\
                                   .filter_by(request_id=reqid)
            if not requester.admin:
                user_requests = user_requests.join(ParametricJobs.request)\
                                             .filter_by(requester_id=requester.id)
            return json.dumps({'data': user_requests.all()}, cls=JSONTableEncoder)

    @staticmethod
    def PUT(jobid, reschedule=False):  # pylint: disable=invalid-name
        """REST Put method."""
        logger.debug("In PUT: reqid = %s, reschedule = %s", reqid, reschedule)
        requester = cherrypy.request.verified_user
        with db_session() as session:
            query = session.query(ParametricJobs).filter_by(id=jobid)
            if not requester.admin:
                query = query.join(ParametricJobs.request)\
                             .filter_by(requester_id=requester.id)
            try:
                job = query.one()
            except NoResultFound:
                logger.error("No ParametricJobs found with id: %s", jobid)
            except MultipleResultsFound:
                logger.error("Multiple ParametricJobs found with id: %s", jobid)
            else:
                if reschedule and not job.reschedule:
                    job.reschedule = True
                    job.status = LOCALSTATUS.Submitting

<commit_msg>adding method docstring.<commit_after>"""ParametricJobs Table."""
import os
import re
import time
import json
import logging
import calendar
from datetime import datetime

import cherrypy
from sqlalchemy import Column, Integer, Boolean, String, PickleType, TIMESTAMP, ForeignKey, Enum
from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy.orm import relationship
from sqlalchemy.orm.exc import NoResultFound, MultipleResultsFound

from rpc.DiracRPCClient import dirac_api_client, ParametricDiracJobClient
from utils.collections import list_splitter
from utils.tempfile_utils import temporary_runscript, temporary_macro
from ..utils import db_session
from ..statuses import LOCALSTATUS
from .SQLTableBase import SQLTableBase
from .JSONTableEncoder import JSONTableEncoder
from .DiracJobs import DiracJobs

logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
UNIXDATE = re.compile(r'(?P<month>[0-9]{2})-(?P<day>[0-9]{2})-(?P<year>[0-9]{4})$')

@cherrypy.expose
class ParametricJobs(SQLTableBase):
    """Jobs SQL Table."""

    __tablename__ = 'parametricjobs'
    id = Column(Integer, primary_key=True)  # pylint: disable=invalid-name
    app = Column(String(250))
    app_version = Column(String(250))
    sim_lfn_outputdir = Column(String(250))
    mctruth_lfn_outputdir = Column(String(250))
    macro = Column(String(250))
    tag = Column(String(250))
    njobs = Column(Integer)
    nevents = Column(Integer)
    seed = Column(Integer)
    fastnest_version = Column(String(250))
    reduction_version = Column(String(250))
    reduction_lfn_inputdir = Column(String(250))
    reduction_lfn_outputdir = Column(String(250))
    der_version = Column(String(250))
    der_lfn_inputdir = Column(String(250))
    der_lfn_outputdir = Column(String(250))
    lzap_version = Column(String(250))
    lzap_lfn_inputdir = Column(String(250))
    lzap_lfn_outputdir = Column(String(250))
    request_id = Column(Integer, ForeignKey('requests.id'), nullable=False)
    request = relationship("Requests", back_populates="parametricjobs")
    status = Column(Enum(LOCALSTATUS), nullable=False)
    reschedule = Column(Boolean, nullable=False, default=False)
    timestamp = Column(TIMESTAMP, nullable=False, default=datetime.utcnow, onupdate=datetime.utcnow)
    num_completed = Column(Integer, nullable=False, default=0)
    num_failed = Column(Integer, nullable=False, default=0)
    num_submitted = Column(Integer, nullable=False, default=0)
    num_running = Column(Integer, nullable=False, default=0)
    diracjobs = relationship("DiracJobs", back_populates="parametricjob")

    @hybrid_property
    def num_other(self):
        """Return the number of jobs in states other than the known ones."""
        return self.njobs - (self.num_submitted + self.num_running + self.num_failed + self.num_completed)

    def submit(self):
        """Submit parametric job."""
        # problems here if not running simulation, there will be no macro so
        # everything including context needs reworking.
#        lfn_root = os.path.join('/lz/user/l/lzproduser.grid.hep.ph.ic.ac.uk', '_'.join(('-'.join((self.app, self.app_version)),
#                                                           '-'.join(('DER', self.der_version)))))

        parametric_job = ParametricDiracJobClient()
        dirac_ids = set()
        if self.app_version is not None:

            macro_name = os.path.splitext(os.path.basename(self.macro or ''))[0]
            livetime_sec_per_beamon = 0.1132698957
            livetimeperjob = str((self.nevents or 1) * livetime_sec_per_beamon)
            unixtime = time.time()
            match = UNIXDATE.search(macro_name)
            if match is not None:
                month, day, year = match.groups()
                unixtime = str(int(calendar.timegm(datetime(int(year), int(month), int(day), 0, 0).utctimetuple())))
            with temporary_runscript(root_version='5.34.32',
                                     root_arch='slc6_gcc44_x86_64',
                                     g4_version='4.9.5.p02',
                                     physics_version='1.4.0',
                                     se='UKI-LT2-IC-HEP-disk',
                                     unixtime=unixtime,
                                     livetimeperjob=livetimeperjob, **self) as runscript,\
                 temporary_macro(self.tag, self.macro or '', self.app, self.nevents) as macro:
                logger.info("Submitting ParametricJob %s, macro: %s to DIRAC", self.id, self.macro)
                for sublist in list_splitter(range(self.seed, self.seed + self.njobs), 1000):
                    with parametric_job as j:
                        j.setName(os.path.splitext(os.path.basename(macro))[0] + '-%(args)s')
                        j.setPlatform('ANY')
                        j.setExecutable(os.path.basename(runscript),
                                        os.path.basename(macro) + ' %(args)s',
                                        'lzproduction_output.log')
                        j.setInputSandbox([runscript, macro])
                        j.setDestination('ANY')
                        j.setParameterSequence('args', sublist, addToWorkflow=False)
                    dirac_ids.update(parametric_job.subjob_ids)
        else:
            with temporary_runscript(root_version='5.34.32',
                                     root_arch='slc6_gcc44_x86_64',
                                     g4_version='4.9.5.p02',
                                     physics_version='1.4.0',
                                     se='UKI-LT2-IC-HEP-disk', **self) as runscript:
                logger.info("Submitting ParametricJob %s, inputdir: %s to DIRAC", self.id, self.reduction_lfn_inputdir or self.der_lfn_inputdir or self.lzap_lfn_inputdir)

                input_lfn_dir=self.reduction_lfn_inputdir or\
                               self.der_lfn_inputdir or \
                               self.lzap_lfn_inputdir
                for sublist in list_splitter(list_lfns(input_lfn_dir), 1000):
                    with parametric_job as j:
                        j.setName("%(args)s")
                        j.setPlatform('ANY')
                        j.setExecutable(os.path.basename(runscript),
                                        '%(args)s',
                                        'lzanalysis_output.log')
                        j.setInputSandbox([runscript])
                        j.setDestination('ANY')
                        j.setParameterSequence('InputData', sublist, addToWorkflow='ParametricInputData')
                        j.setParameterSequence('args',
                                               [os.path.basename(l) for l in sublist],
                                               addToWorkflow=False)
                    dirac_ids.update(parametric_job.subjob_ids)

        with db_session() as session:
            session.bulk_insert_mappings(DiracJobs, [{'id': i, 'parametricjob_id': self.id}
                                                     for i in parametric_job.subjob_ids])

    def reset(self):
        """Reset parametric job."""
        with db_session(reraise=False) as session:
            dirac_jobs = session.query(DiracJobs).filter_by(parametricjob_id=self.id)
            dirac_job_ids = [j.id for j in dirac_jobs.all()]
            dirac_jobs.delete(synchronize_session=False)
        with dirac_api_client() as dirac:
            logger.info("Removing Dirac jobs %s from ParametricJob %s", dirac_job_ids, self.id)
            dirac.kill(dirac_job_ids)
            dirac.delete(dirac_job_ids)

    def delete_dirac_jobs(self, session):
        """Delete associated DIRAC jobs."""
        logger.info("Deleting DiracJobs for ParametricJob id: %s", self.id)
        session.query(DiracJobs)\
               .filter_by(parametricjob_id=self.id)\
               .delete(synchronize_session=False)

    def update_status(self):
        """Update the status of parametric job."""
        local_statuses = DiracJobs.update_status(self)
        # could just have DiracJobs return this... maybe better
#        local_statuses = Counter(status.local_status for status in dirac_statuses.elements())
        status = max(local_statuses or [self.status])
        with db_session() as session:
            this = session.merge(self)
            this.status = status
            this.num_completed = local_statuses['Completed']
            this.num_failed = local_statuses['Failed']
            this.num_submitted = local_statuses['Submitted']
            this.num_running = local_statuses['Running']
            this.reschedule = False
        return status

    @staticmethod
    def GET(reqid):  # pylint: disable=invalid-name
        """
        REST Get method.

        Returns all ParametricJobs for a given request id.
        """
        logger.debug("In GET: reqid = %s", reqid)
        requester = cherrypy.request.verified_user
        with db_session() as session:
            user_requests = session.query(ParametricJobs)\
                                   .filter_by(request_id=reqid)
            if not requester.admin:
                user_requests = user_requests.join(ParametricJobs.request)\
                                             .filter_by(requester_id=requester.id)
            return json.dumps({'data': user_requests.all()}, cls=JSONTableEncoder)

    @staticmethod
    def PUT(jobid, reschedule=False):  # pylint: disable=invalid-name
        """REST Put method."""
        logger.debug("In PUT: reqid = %s, reschedule = %s", reqid, reschedule)
        requester = cherrypy.request.verified_user
        with db_session() as session:
            query = session.query(ParametricJobs).filter_by(id=jobid)
            if not requester.admin:
                query = query.join(ParametricJobs.request)\
                             .filter_by(requester_id=requester.id)
            try:
                job = query.one()
            except NoResultFound:
                logger.error("No ParametricJobs found with id: %s", jobid)
            except MultipleResultsFound:
                logger.error("Multiple ParametricJobs found with id: %s", jobid)
            else:
                if reschedule and not job.reschedule:
                    job.reschedule = True
                    job.status = LOCALSTATUS.Submitting

<|endoftext|>


<commit_before># Copyright 2019 The Bazel Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""PostCSS runner rule.

This generates the internal PostCSS runner as a nodejs_binary, that can later
be used as an executable."""

load(":runner_bin.bzl", "postcss_runner_bin")

ERROR_INPUT_NO_PLUGINS = "No plugins were provided"

def _postcss_runner_src_impl(ctx):
    if len(ctx.attr.plugins.items()) == 0:
        fail(ERROR_INPUT_NO_PLUGINS)

    plugins = []
    for plugin_key, plugin_options in ctx.attr.plugins.items():
        plugins.append("require('%s').apply(this, %s)" %
                       (plugin_key, plugin_options if plugin_options else "[]"))

    ctx.actions.expand_template(
        template = ctx.file.template,
        output = ctx.outputs.postcss_runner_src,
        substitutions = {
            # The array of PostCSS plugin objects.
            "TEMPLATED_plugins": "[%s]" % (",".join(plugins)) if plugins else "",
            # Lowercase Python boolean converts to true/false.
            "TEMPLATED_map_annotation": str(ctx.attr.map_annotation).lower(),
        },
    )

postcss_runner_src = rule(
    implementation = _postcss_runner_src_impl,
    attrs = {
        "plugins": attr.string_dict(
            mandatory = True,
        ),
        "map_annotation": attr.bool(
            default = False,
        ),
        "template": attr.label(
            allow_single_file = [".js"],
            mandatory = True,
        ),
    },
    outputs = {
        "postcss_runner_src": "%{name}.js",
    },
)

def postcss_gen_runner(
        name,
        plugins,
        deps,
        map_annotation,
        visibility = None):
    """Generates a PostCSS runner binary.

    Pass the result to a postcss_run rule to apply the runner to a css file.

    Args:
        name: The name of the build rule.
        plugins: A map of plugin Node.js require paths (following the
            requirements of rules_nodejs), with values being config objects
            for each respective plugin.
        deps: A list of NodeJS modules the config depends on. The postcss module
            is included by default.
        map_annotation: Whether to add (or modify, if already existing) the
            sourceMappingURL comment in the output .css to point to the output
            .css.map.
        visibility: The visibility of the build rule.
    """

    runner_src_name = "%s.runner_src" % name

    postcss_runner_src(
        name = runner_src_name,
        plugins = plugins,
        map_annotation = map_annotation,
        template = "//internal:runner-template.js",
        visibility = ["//visibility:private"],
    )

    postcss_runner_bin(
        name = name,
        src = runner_src_name,
        deps = [
            "@npm//minimist",
            "@npm//postcss",
        ] + deps,
        visibility = visibility,
    )
<commit_msg>Use absolute path to refer to PostCSS runner<commit_after># Copyright 2019 The Bazel Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""PostCSS runner rule.

This generates the internal PostCSS runner as a nodejs_binary, that can later
be used as an executable."""

load(":runner_bin.bzl", "postcss_runner_bin")

ERROR_INPUT_NO_PLUGINS = "No plugins were provided"

def _postcss_runner_src_impl(ctx):
    if len(ctx.attr.plugins.items()) == 0:
        fail(ERROR_INPUT_NO_PLUGINS)

    plugins = []
    for plugin_key, plugin_options in ctx.attr.plugins.items():
        plugins.append("require('%s').apply(this, %s)" %
                       (plugin_key, plugin_options if plugin_options else "[]"))

    ctx.actions.expand_template(
        template = ctx.file.template,
        output = ctx.outputs.postcss_runner_src,
        substitutions = {
            # The array of PostCSS plugin objects.
            "TEMPLATED_plugins": "[%s]" % (",".join(plugins)) if plugins else "",
            # Lowercase Python boolean converts to true/false.
            "TEMPLATED_map_annotation": str(ctx.attr.map_annotation).lower(),
        },
    )

postcss_runner_src = rule(
    implementation = _postcss_runner_src_impl,
    attrs = {
        "plugins": attr.string_dict(
            mandatory = True,
        ),
        "map_annotation": attr.bool(
            default = False,
        ),
        "template": attr.label(
            allow_single_file = [".js"],
            mandatory = True,
        ),
    },
    outputs = {
        "postcss_runner_src": "%{name}.js",
    },
)

def postcss_gen_runner(
        name,
        plugins,
        deps,
        map_annotation,
        visibility = None):
    """Generates a PostCSS runner binary.

    Pass the result to a postcss_run rule to apply the runner to a css file.

    Args:
        name: The name of the build rule.
        plugins: A map of plugin Node.js require paths (following the
            requirements of rules_nodejs), with values being config objects
            for each respective plugin.
        deps: A list of NodeJS modules the config depends on. The postcss module
            is included by default.
        map_annotation: Whether to add (or modify, if already existing) the
            sourceMappingURL comment in the output .css to point to the output
            .css.map.
        visibility: The visibility of the build rule.
    """

    runner_src_name = "%s.runner_src" % name

    postcss_runner_src(
        name = runner_src_name,
        plugins = plugins,
        map_annotation = map_annotation,
        template = "@build_bazel_rules_postcss//internal:runner-template.js",
        visibility = ["//visibility:private"],
    )

    postcss_runner_bin(
        name = name,
        src = runner_src_name,
        deps = [
            "@npm//minimist",
            "@npm//postcss",
        ] + deps,
        visibility = visibility,
    )
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# FOGLAMP_BEGIN
# See: http://foglamp.readthedocs.io/
# FOGLAMP_END

import asyncio
import datetime
import os
import time
import uuid

import pytest

from foglamp.core.scheduler import (Scheduler, IntervalSchedule, ScheduleNotFoundError, Task,
                                    Schedule, TimedSchedule, ManualSchedule, StartUpSchedule)

__author__ = "Terris Linenbach"
__copyright__ = "Copyright (c) 2017 OSIsoft, LLC"
__license__ = "Apache 2.0"
__version__ = "${VERSION}"

class TestScheduler:
    @staticmethod
    async def stop_scheduler(scheduler: Scheduler)->None:
        while True:
            try:
                await scheduler.stop()
                break
            except TimeoutError:
                await asyncio.sleep(1)

    @pytest.mark.asyncio
    async def test_stop(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.exclusive = False
        interval_schedule.name = 'sleep1'
        interval_schedule.process_name = "sleep1"
        interval_schedule.repeat = datetime.timedelta(seconds=1)

        await scheduler.save_schedule(interval_schedule)
        await asyncio.sleep(10)

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_create_interval(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        assert interval_schedule.schedule_type == Schedule.Type.INTERVAL

        interval_schedule.name = 'sleep10'
        interval_schedule.process_name = "sleep10"
        interval_schedule.repeat = datetime.timedelta(seconds=1)

        await scheduler.save_schedule(interval_schedule)

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_update(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.name = 'sleep10'
        interval_schedule.process_name = "sleep10"

        await scheduler.save_schedule(interval_schedule)

        tasks = await scheduler.get_running_tasks()
        await asyncio.sleep(1)
        assert len(tasks) == 0

        interval_schedule.name = 'updated'
        interval_schedule.process_name = "sleep1"
        interval_schedule.repeat = datetime.timedelta(seconds=5)

        await scheduler.save_schedule(interval_schedule)
        await asyncio.sleep(6)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 1

        interval_schedule.exclusive = False
        await scheduler.save_schedule(interval_schedule)

        # Check able to get same schedule after restart
        # Check fields have been modified
        await self.stop_scheduler(scheduler)
        scheduler = Scheduler()
        await scheduler.start()

        schedule = await scheduler.get_schedule(interval_schedule.schedule_id)
        assert schedule.process_name == 'sleep1'
        assert schedule.name == 'updated'
        assert schedule.repeat.seconds == 5
        assert not schedule.exclusive

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_startup_schedule(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        startup_schedule = StartUpSchedule()
        startup_schedule.name = 'startup schedule'
        startup_schedule.process_name = 'sleep30'
        startup_schedule.repeat = datetime.timedelta(seconds=0)

        await scheduler.save_schedule(startup_schedule)

        await asyncio.sleep(1)
        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 0

        await scheduler.get_schedule(startup_schedule.schedule_id)

        await self.stop_scheduler(scheduler)

        scheduler = Scheduler()
        await scheduler.start()

        await asyncio.sleep(2)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 1

        scheduler.max_running_tasks = 0
        await scheduler.cancel_task(tasks[0].task_id)

        await asyncio.sleep(2)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 0

        scheduler.max_running_tasks = 1

        await asyncio.sleep(2)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 1

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_manual_schedule(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        manual_schedule = ManualSchedule()
        manual_schedule.name = 'manual task'
        manual_schedule.process_name = 'sleep10'

        await scheduler.save_schedule(manual_schedule)
        manual_schedule = await scheduler.get_schedule(manual_schedule.schedule_id)

        await scheduler.queue_task(manual_schedule.schedule_id)
        await asyncio.sleep(5)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 1

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_max_processes(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        # 2 maximum tasks

        # 1 runs at 1 second
        # 2 runs at 2 seconds
        # 3 runs at 11 seconds
        # 4 runs at 12 seconds
        # 5 runs at 21 seconds
        # 6 runs at 22 seconds
        # 7 runs at 31 seconds
        # 8 runs at 32 seconds
        # Total: 6

        scheduler.max_running_tasks = 2

        interval_schedule = IntervalSchedule()
        interval_schedule.repeat = datetime.timedelta(seconds=1)
        interval_schedule.name = 'max active'
        interval_schedule.exclusive = False
        interval_schedule.process_name = 'sleep10'

        await scheduler.save_schedule(interval_schedule)

        await asyncio.sleep(30.3)
        scheduler.max_running_tasks = 0

        tasks = await scheduler.get_tasks(10)
        assert len(tasks) == 6

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 2

        # They end...
        await asyncio.sleep(20)

        scheduler.max_running_tasks = 10

        await asyncio.sleep(11)
        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 10

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_timed_schedule(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        timed_schedule = TimedSchedule()

        # Tuesday August 8 2017 8:00:00 AM PDT
        now = 1502204400
        scheduler.current_time = now

        timed_schedule.name = 'timed'
        timed_schedule.process_name = 'sleep10'
        timed_schedule.day = 2
        timed_schedule.time = datetime.time(hour=8)

        os.environ["TZ"] = "PST8PDT"
        time.tzset()

        await scheduler.save_schedule(timed_schedule)
        await asyncio.sleep(1)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 1

        timed_schedule = await scheduler.get_schedule(
                            uuid.UUID(str(timed_schedule.schedule_id)))

        assert timed_schedule.time.hour == 8
        assert timed_schedule.time.minute == 0
        assert timed_schedule.time.second == 0
        assert timed_schedule.day == 2

        del os.environ["TZ"]
        time.tzset()

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_delete(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.name = 'deletetest'
        interval_schedule.process_name = "sleep1"
        await scheduler.save_schedule(interval_schedule)

        await asyncio.sleep(5)

        await scheduler.delete_schedule(interval_schedule.schedule_id)

        try:
            await scheduler.delete_schedule(interval_schedule.schedule_id)
            assert False
        except ScheduleNotFoundError:
            pass

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_cancel(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.name = 'cancel_test'
        interval_schedule.process_name = 'sleep30'
        await scheduler.save_schedule(interval_schedule)

        await asyncio.sleep(5)
        tasks = await scheduler.get_running_tasks()
        await scheduler.cancel_task(tasks[0].task_id)

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_get_schedule(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.name = 'get_schedule_test'
        interval_schedule.process_name = "sleep30"
        await scheduler.save_schedule(interval_schedule)

        schedules = await scheduler.get_schedules()
        assert len(schedules) == 1

        await scheduler.get_schedule(interval_schedule.schedule_id)

        try:
            await scheduler.get_schedule(uuid.uuid4())
            assert False
        except ScheduleNotFoundError:
            pass

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_get_task(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.name = 'get_task'
        interval_schedule.process_name = "sleep30"
        await scheduler.save_schedule(interval_schedule)
        await asyncio.sleep(1)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks)

        task = await scheduler.get_task(tasks[0].task_id)
        assert task

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_get_tasks(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.name = 'get_tasks'
        interval_schedule.process_name = "sleep5"
        interval_schedule.repeat = datetime.timedelta(seconds=1)
        interval_schedule.exclusive = False
        await scheduler.save_schedule(interval_schedule)

        await asyncio.sleep(15)

        tasks = await scheduler.get_tasks(50)
        assert len(tasks) > 1
        assert tasks[0].state == Task.State.RUNNING
        assert tasks[-1].state == Task.State.COMPLETE

        assert len(tasks)
        tasks2 = await scheduler.get_tasks(1)
        assert len(tasks2) == 1
        assert tasks[0].start_time == tasks2[0].start_time

        tasks = await scheduler.get_tasks(
            where=Task.attr.state.in_(Task.State.RUNNING),
            sort=[Task.attr.state.desc], offset=1)
        """
        tasks = await scheduler.get_tasks(
            where=(Task.attr.state == Task.State.RUNNING).or_(
                Task.attr.state == Task.State.RUNNING,
                Task.attr.state == Task.State.RUNNING).and_(
                Task.attr.state.in_(int(Task.State.RUNNING)),
                Task.attr.state.in_(int(Task.State.RUNNING)).or_(
                    Task.attr.state.in_(int(Task.State.RUNNING)))),
            sort=Task.attr.state.desc)
        """
        assert len(tasks)

        await self.stop_scheduler(scheduler)
<commit_msg>Almost<commit_after># -*- coding: utf-8 -*-

# FOGLAMP_BEGIN
# See: http://foglamp.readthedocs.io/
# FOGLAMP_END

import asyncio
import datetime
import os
import time
import uuid

import pytest

from foglamp.core.scheduler import (Scheduler, IntervalSchedule, ScheduleNotFoundError, Task,
                                    Schedule, TimedSchedule, ManualSchedule, StartUpSchedule)

__author__ = "Terris Linenbach"
__copyright__ = "Copyright (c) 2017 OSIsoft, LLC"
__license__ = "Apache 2.0"
__version__ = "${VERSION}"

class TestScheduler:
    @staticmethod
    async def stop_scheduler(scheduler: Scheduler)->None:
        while True:
            try:
                await scheduler.stop()
                break
            except TimeoutError:
                await asyncio.sleep(1)

    @pytest.mark.asyncio
    async def test_stop(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.exclusive = False
        interval_schedule.name = 'sleep1'
        interval_schedule.process_name = "sleep1"
        interval_schedule.repeat = datetime.timedelta(seconds=1)

        await scheduler.save_schedule(interval_schedule)
        await asyncio.sleep(10)

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_create_interval(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        assert interval_schedule.schedule_type == Schedule.Type.INTERVAL

        interval_schedule.name = 'sleep10'
        interval_schedule.process_name = "sleep10"
        interval_schedule.repeat = datetime.timedelta(seconds=1)

        await scheduler.save_schedule(interval_schedule)

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_update(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.name = 'sleep10'
        interval_schedule.process_name = "sleep10"

        await scheduler.save_schedule(interval_schedule)

        tasks = await scheduler.get_running_tasks()
        await asyncio.sleep(1)
        assert len(tasks) == 0

        interval_schedule.name = 'updated'
        interval_schedule.process_name = "sleep1"
        interval_schedule.repeat = datetime.timedelta(seconds=5)

        await scheduler.save_schedule(interval_schedule)
        await asyncio.sleep(6)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 1

        interval_schedule.exclusive = False
        await scheduler.save_schedule(interval_schedule)

        # Check able to get same schedule after restart
        # Check fields have been modified
        await self.stop_scheduler(scheduler)
        scheduler = Scheduler()
        await scheduler.start()

        schedule = await scheduler.get_schedule(interval_schedule.schedule_id)
        assert schedule.process_name == 'sleep1'
        assert schedule.name == 'updated'
        assert schedule.repeat.seconds == 5
        assert not schedule.exclusive

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_startup_schedule(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        startup_schedule = StartUpSchedule()
        startup_schedule.name = 'startup schedule'
        startup_schedule.process_name = 'sleep30'
        startup_schedule.repeat = datetime.timedelta(seconds=0)

        await scheduler.save_schedule(startup_schedule)

        await asyncio.sleep(1)
        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 0

        await scheduler.get_schedule(startup_schedule.schedule_id)

        await self.stop_scheduler(scheduler)

        scheduler = Scheduler()
        await scheduler.start()

        await asyncio.sleep(2)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 1

        scheduler.max_running_tasks = 0
        await scheduler.cancel_task(tasks[0].task_id)

        await asyncio.sleep(2)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 0

        scheduler.max_running_tasks = 1

        await asyncio.sleep(2)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 1

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_manual_schedule(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        manual_schedule = ManualSchedule()
        manual_schedule.name = 'manual task'
        manual_schedule.process_name = 'sleep10'

        await scheduler.save_schedule(manual_schedule)
        manual_schedule = await scheduler.get_schedule(manual_schedule.schedule_id)

        await scheduler.queue_task(manual_schedule.schedule_id)
        await asyncio.sleep(5)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 1

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_max_processes(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        # 2 maximum tasks

        # 1 runs at 1 second
        # 2 runs at 2 seconds
        # 3 runs at 11 seconds
        # 4 runs at 12 seconds
        # 5 runs at 21 seconds
        # 6 runs at 22 seconds
        # 7 runs at 31 seconds
        # 8 runs at 32 seconds
        # Total: 6

        scheduler.max_running_tasks = 2

        interval_schedule = IntervalSchedule()
        interval_schedule.repeat = datetime.timedelta(seconds=1)
        interval_schedule.name = 'max active'
        interval_schedule.exclusive = False
        interval_schedule.process_name = 'sleep10'

        await scheduler.save_schedule(interval_schedule)

        await asyncio.sleep(30.3)
        scheduler.max_running_tasks = 0

        tasks = await scheduler.get_tasks(10)
        assert len(tasks) == 6

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 2

        # They end...
        await asyncio.sleep(20)

        scheduler.max_running_tasks = 10

        await asyncio.sleep(11)
        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 10

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_timed_schedule(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        timed_schedule = TimedSchedule()

        # Tuesday August 8 2017 8:00:00 AM PDT
        now = 1502204400
        scheduler.current_time = now

        timed_schedule.name = 'timed'
        timed_schedule.process_name = 'sleep10'
        timed_schedule.day = 2
        timed_schedule.time = datetime.time(hour=8)

        os.environ["TZ"] = "PST8PDT"
        time.tzset()

        await scheduler.save_schedule(timed_schedule)
        await asyncio.sleep(1)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks) == 1

        timed_schedule = await scheduler.get_schedule(
                            uuid.UUID(str(timed_schedule.schedule_id)))

        assert timed_schedule.time.hour == 8
        assert timed_schedule.time.minute == 0
        assert timed_schedule.time.second == 0
        assert timed_schedule.day == 2

        del os.environ["TZ"]
        time.tzset()

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_delete(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.name = 'deletetest'
        interval_schedule.process_name = "sleep1"
        await scheduler.save_schedule(interval_schedule)

        await asyncio.sleep(5)

        await scheduler.delete_schedule(interval_schedule.schedule_id)

        try:
            await scheduler.delete_schedule(interval_schedule.schedule_id)
            assert False
        except ScheduleNotFoundError:
            pass

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_cancel(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.name = 'cancel_test'
        interval_schedule.process_name = 'sleep30'
        await scheduler.save_schedule(interval_schedule)

        await asyncio.sleep(5)
        tasks = await scheduler.get_running_tasks()
        await scheduler.cancel_task(tasks[0].task_id)

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_get_schedule(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.name = 'get_schedule_test'
        interval_schedule.process_name = "sleep30"
        await scheduler.save_schedule(interval_schedule)

        schedules = await scheduler.get_schedules()
        assert len(schedules) == 1

        await scheduler.get_schedule(interval_schedule.schedule_id)

        try:
            await scheduler.get_schedule(uuid.uuid4())
            assert False
        except ScheduleNotFoundError:
            pass

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_get_task(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.name = 'get_task'
        interval_schedule.process_name = "sleep30"
        await scheduler.save_schedule(interval_schedule)
        await asyncio.sleep(1)

        tasks = await scheduler.get_running_tasks()
        assert len(tasks)

        task = await scheduler.get_task(tasks[0].task_id)
        assert task

        await self.stop_scheduler(scheduler)

    @pytest.mark.asyncio
    async def test_get_tasks(self):
        scheduler = Scheduler()

        await scheduler.populate_test_data()
        await scheduler.start()

        interval_schedule = IntervalSchedule()
        interval_schedule.name = 'get_tasks'
        interval_schedule.process_name = "sleep5"
        interval_schedule.repeat = datetime.timedelta(seconds=1)
        interval_schedule.exclusive = False
        await scheduler.save_schedule(interval_schedule)

        await asyncio.sleep(15)

        tasks = await scheduler.get_tasks(50)
        assert len(tasks) > 1
        assert tasks[0].state == Task.State.RUNNING
        assert tasks[-1].state == Task.State.COMPLETE

        tasks = await scheduler.get_tasks(1)
        assert len(tasks) == 1

        tasks = await scheduler.get_tasks(
            where=Task.attr.state.in_(int(Task.State.RUNNING)),
            sort=[Task.attr.state.desc], offset=50)
        assert len(tasks) == 0

        tasks = await scheduler.get_tasks(
            where=(Task.attr.state == Task.State.RUNNING).or_(
                Task.attr.state == Task.State.RUNNING,
                Task.attr.state == Task.State.RUNNING).and_(
                Task.attr.state.in_(int(Task.State.RUNNING)),
                Task.attr.state.in_(int(Task.State.RUNNING)).or_(
                    Task.attr.state.in_(int(Task.State.RUNNING)))),
            sort=Task.attr.state.desc)
        assert len(tasks)

        await self.stop_scheduler(scheduler)
<|endoftext|>


<commit_before>"""
Copyright 2018 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from abc import ABCMeta
from abc import abstractmethod
from six.moves import zip
import numpy as np
import six
import tensorflow as tf
from google.protobuf import text_format

class ModelWrapper(six.with_metaclass(ABCMeta, object)):
  """Simple wrapper of the for models with session object for TCAV.

    Supports easy inference with no need to deal with the feed_dicts.
  """

  @abstractmethod
  def __init__(self, model_path=None, node_dict=None):
    """Initialize the wrapper.

    Optionally create a session, load
    the model from model_path to this session, and map the
    input/output and bottleneck tensors.

    Args:
      model_path: one of the following: 1) Directory path to checkpoint 2)
        Directory path to SavedModel 3) File path to frozen graph.pb 4) File
        path to frozen graph.pbtxt
      node_dict: mapping from a short name to full input/output and bottleneck
        tensor names. Users should pass 'input' and 'prediction'
        as keys and the corresponding input and prediction tensor
        names as values in node_dict. Users can additionally pass bottleneck
        tensor names for which gradient Ops will be added later.
    """
    # A dictionary of bottleneck tensors.
    self.bottlenecks_tensors = None
    # A dictionary of input, 'logit' and prediction tensors.
    self.ends = None
    # The model name string.
    self.model_name = None
    # a place holder for index of the neuron/class of interest.
    # usually defined under the graph. For example:
    # with g.as_default():
    #   self.tf.placeholder(tf.int64, shape=[None])
    self.y_input = None
    # The tensor representing the loss (used to calculate derivative).
    self.loss = None
    # If tensors in the loaded graph are prefixed with 'import/'
    self.import_prefix = False

    if model_path:
      self._try_loading_model(model_path)
    if node_dict:
      self._find_ends_and_bottleneck_tensors(node_dict)

  def _try_loading_model(self, model_path):
    """ Load model from model_path.

    TF models are often saved in one of the three major formats:
      1) Checkpoints with ckpt.meta, ckpt.data, and ckpt.index.
      2) SavedModel format with saved_model.pb and variables/.
      3) Frozen graph in .pb or .pbtxt format.
    When model_path is specified, model is loaded in one of the
    three formats depending on the model_path. When model_path is
    ommitted, child wrapper is responsible for loading the model.
    """
    try:
      self.sess = tf.Session(graph=tf.Graph())
      with self.sess.graph.as_default():
        if tf.io.gfile.isdir(model_path):
          ckpt = tf.train.latest_checkpoint(model_path)
          if ckpt:
            tf.logging.info('Loading from the latest checkpoint.')
            saver = tf.train.import_meta_graph(ckpt + '.meta')
            saver.restore(self.sess, ckpt)
          else:
            tf.logging.info('Loading from SavedModel dir.')
            tf.saved_model.loader.load(self.sess, ['serve'], model_path)
        else:
          input_graph_def = tf.GraphDef()
          if model_path.endswith('.pb'):
            tf.logging.info('Loading from frozen binary graph.')
            with tf.io.gfile.GFile(model_path, 'rb') as f:
              input_graph_def.ParseFromString(f.read())
          else:
            tf.logging.info('Loading from frozen text graph.')
            with tf.io.gfile.GFile(model_path) as f:
              text_format.Parse(f.read(), input_graph_def)
          tf.import_graph_def(input_graph_def)
          self.import_prefix = True

    except Exception as e:
      template = 'An exception of type {0} occurred ' \
                 'when trying to load model from {1}. ' \
                 'Arguments:\n{2!r}'
      tf.logging.warn(template.format(type(e).__name__, model_path, e.args))

  def _find_ends_and_bottleneck_tensors(self, node_dict):
    """ Find tensors from the graph by their names.

    Depending on how the model is loaded, tensors in the graph
    may or may not have 'import/' prefix added to every tensor name.
    This is true even if the tensors already have 'import/' prefix.
    The 'ends' and 'bottlenecks_tensors' dictionary should map to tensors
    with the according name.
    """
    self.bottlenecks_tensors = {}
    self.ends = {}
    for k, v in node_dict.iteritems():
      if self.import_prefix:
        v = 'import/' + v
      tensor = self.sess.graph.get_operation_by_name(v.strip(':0')).outputs[0]
      if k == 'input' or k == 'prediction':
        self.ends[k] = tensor
      else:
        self.bottlenecks_tensors[k] = tensor

  def _make_gradient_tensors(self):
    """Makes gradient tensors for all bottleneck tensors.
    """
    self.bottlenecks_gradients = {}
    for bn in self.bottlenecks_tensors:
      self.bottlenecks_gradients[bn] = tf.gradients(
          self.loss, self.bottlenecks_tensors[bn])[0]

  def get_gradient(self, acts, y, bottleneck_name, example):
    """Return the gradient of the loss with respect to the bottleneck_name.

    Args:
      acts: activation of the bottleneck
      y: index of the logit layer
      bottleneck_name: name of the bottleneck to get gradient wrt.
      example: input example. Unused by default. Necessary for getting gradients
        from certain models, such as BERT.

    Returns:
      the gradient array.
    """
    return self.sess.run(self.bottlenecks_gradients[bottleneck_name], {
        self.bottlenecks_tensors[bottleneck_name]: acts,
        self.y_input: y
    })

  def get_predictions(self, examples):
    """Get prediction of the examples.

    Args:
      imgs: array of examples to get predictions

    Returns:
      array of predictions
    """
    return self.adjust_prediction(
        self.sess.run(self.ends['prediction'], {self.ends['input']: examples}))

  def adjust_prediction(self, pred_t):
    """Adjust the prediction tensor to be the expected shape.

    Defaults to a no-op, but necessary to override for GoogleNet
    Returns:
      pred_t: pred_tensor.
    """
    return pred_t

  def reshape_activations(self, layer_acts):
    """Reshapes layer activations as needed to feed through the model network.

    Override this for models that require reshaping of the activations for use
    in TCAV.

    Args:
      layer_acts: Activations as returned by run_examples.

    Returns:
      Activations in model-dependent form; the default is a squeezed array (i.e.
      at most one dimensions of size 1).
    """
    return np.asarray(layer_acts).squeeze()

  def label_to_id(self, label):
    """Convert label (string) to index in the logit layer (id).

    Override this method if label to id mapping is known. Otherwise,
    default id 0 is used.
    """
    tf.logging.warn('label_to_id undefined. Defaults to returning 0.')
    return 0

  def id_to_label(self, idx):
    """Convert index in the logit layer (id) to label (string).

    Override this method if id to label mapping is known.
    """
    return str(idx)

  def run_examples(self, examples, bottleneck_name):
    """Get activations at a bottleneck for provided examples.

    Args:
      examples: example data to feed into network.
      bottleneck_name: string, should be key of self.bottlenecks_tensors

    Returns:
      Activations in the given layer.
    """
    return self.sess.run(self.bottlenecks_tensors[bottleneck_name],
                         {self.ends['input']: examples})

class ImageModelWrapper(ModelWrapper):
  """Wrapper base class for image models."""

  def __init__(self, image_shape):
    super(ModelWrapper, self).__init__()
    # shape of the input image in this model
    self.image_shape = image_shape

  def get_image_shape(self):
    """returns the shape of an input image."""
    return self.image_shape

class PublicImageModelWrapper(ImageModelWrapper):
  """Simple wrapper of the public image models with session object.
  """
  def __init__(self,
               sess,
               model_fn_path,
               labels_path,
               image_shape,
               endpoints_dict,
               scope):
    super(PublicImageModelWrapper, self).__init__(image_shape)
    self.labels = tf.io.gfile.GFile(labels_path).read().splitlines()
    self.ends = PublicImageModelWrapper.import_graph(model_fn_path,
                                                     endpoints_dict,
                                                     self.image_value_range,
                                                     scope=scope)
    self.bottlenecks_tensors = PublicImageModelWrapper.get_bottleneck_tensors(
        scope)
    graph = tf.get_default_graph()

    # Construct gradient ops.
    with graph.as_default():
      self.y_input = tf.placeholder(tf.int64, shape=[None])

      self.pred = tf.expand_dims(self.ends['prediction'][0], 0)
      self.loss = tf.reduce_mean(
          tf.nn.softmax_cross_entropy_with_logits(
              labels=tf.one_hot(
                  self.y_input,
                  self.ends['prediction'].get_shape().as_list()[1]),
              logits=self.pred))
    self._make_gradient_tensors()

  def id_to_label(self, idx):
    return self.labels[idx]

  def label_to_id(self, label):
    return self.labels.index(label)

  @staticmethod
  def create_input(t_input, image_value_range):
    """Create input tensor."""
    def forget_xy(t):
      """Forget sizes of dimensions [1, 2] of a 4d tensor."""
      zero = tf.identity(0)
      return t[:, zero:, zero:, :]

    t_prep_input = t_input
    if len(t_prep_input.shape) == 3:
      t_prep_input = tf.expand_dims(t_prep_input, 0)
    t_prep_input = forget_xy(t_prep_input)
    lo, hi = image_value_range
    t_prep_input = lo + t_prep_input * (hi-lo)
    return t_input, t_prep_input

  # From Alex's code.
  @staticmethod
  def get_bottleneck_tensors(scope):
    """Add Inception bottlenecks and their pre-Relu versions to endpoints dict."""
    graph = tf.get_default_graph()
    bn_endpoints = {}
    for op in graph.get_operations():
      if op.name.startswith(scope+'/') and 'Concat' in op.type:
        name = op.name.split('/')[1]
        bn_endpoints[name] = op.outputs[0]
    return bn_endpoints

  # Load graph and import into graph used by our session
  @staticmethod
  def import_graph(saved_path, endpoints, image_value_range, scope='import'):
    t_input = tf.placeholder(np.float32, [None, None, None, 3])
    graph = tf.Graph()
    assert graph.unique_name(scope, False) == scope, (
        'Scope "%s" already exists. Provide explicit scope names when '
        'importing multiple instances of the model.') % scope

    graph_def = tf.GraphDef.FromString(tf.io.gfile.GFile(saved_path, 'rb').read())

    with tf.name_scope(scope) as sc:
      t_input, t_prep_input = PublicImageModelWrapper.create_input(
          t_input, image_value_range)

      graph_inputs = {}
      graph_inputs[endpoints['input']] = t_prep_input
      myendpoints = tf.import_graph_def(
          graph_def, graph_inputs, list(endpoints.values()), name=sc)
      myendpoints = dict(list(zip(list(endpoints.keys()), myendpoints)))
      myendpoints['input'] = t_input
    return myendpoints

class GoogleNetWrapper_public(PublicImageModelWrapper):

  def __init__(self, sess, model_saved_path, labels_path):
    image_shape_v1 = [224, 224, 3]
    self.image_value_range = (-117, 255-117)
    endpoints_v1 = dict(
        input='input:0',
        logit='softmax2_pre_activation:0',
        prediction='output2:0',
        pre_avgpool='mixed5b:0',
        logit_weight='softmax2_w:0',
        logit_bias='softmax2_b:0',
    )
    self.sess = sess
    super(GoogleNetWrapper_public, self).__init__(sess,
                                                  model_saved_path,
                                                  labels_path,
                                                  image_shape_v1,
                                                  endpoints_v1,
                                                  scope='v1')
    self.model_name = 'GoogleNet_public'

  def adjust_prediction(self, pred_t):
    # Each pred outputs 16, 1008 matrix. The prediction value is the first row.
    # Following tfzoo convention.
    return pred_t[::16]

class InceptionV3Wrapper_public(PublicImageModelWrapper):
  def __init__(self, sess, model_saved_path, labels_path):
    self.image_value_range = (-1, 1)
    image_shape_v3 = [299, 299, 3]
    endpoints_v3 = dict(
        input='Mul:0',
        logit='softmax/logits:0',
        prediction='softmax:0',
        pre_avgpool='mixed_10/join:0',
        logit_weight='softmax/weights:0',
        logit_bias='softmax/biases:0',
    )

    self.sess = sess
    super(InceptionV3Wrapper_public, self).__init__(sess,
                                                  model_saved_path,
                                                  labels_path,
                                                  image_shape_v3,
                                                  endpoints_v3,
                                                  scope='v3')
    self.model_name = 'InceptionV3_public'
<commit_msg>Updating model.py to use six<commit_after>"""
Copyright 2018 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from abc import ABCMeta
from abc import abstractmethod
from six.moves import zip
import numpy as np
import six
import tensorflow as tf
from google.protobuf import text_format

class ModelWrapper(six.with_metaclass(ABCMeta, object)):
  """Simple wrapper of the for models with session object for TCAV.

    Supports easy inference with no need to deal with the feed_dicts.
  """

  @abstractmethod
  def __init__(self, model_path=None, node_dict=None):
    """Initialize the wrapper.

    Optionally create a session, load
    the model from model_path to this session, and map the
    input/output and bottleneck tensors.

    Args:
      model_path: one of the following: 1) Directory path to checkpoint 2)
        Directory path to SavedModel 3) File path to frozen graph.pb 4) File
        path to frozen graph.pbtxt
      node_dict: mapping from a short name to full input/output and bottleneck
        tensor names. Users should pass 'input' and 'prediction'
        as keys and the corresponding input and prediction tensor
        names as values in node_dict. Users can additionally pass bottleneck
        tensor names for which gradient Ops will be added later.
    """
    # A dictionary of bottleneck tensors.
    self.bottlenecks_tensors = None
    # A dictionary of input, 'logit' and prediction tensors.
    self.ends = None
    # The model name string.
    self.model_name = None
    # a place holder for index of the neuron/class of interest.
    # usually defined under the graph. For example:
    # with g.as_default():
    #   self.tf.placeholder(tf.int64, shape=[None])
    self.y_input = None
    # The tensor representing the loss (used to calculate derivative).
    self.loss = None
    # If tensors in the loaded graph are prefixed with 'import/'
    self.import_prefix = False

    if model_path:
      self._try_loading_model(model_path)
    if node_dict:
      self._find_ends_and_bottleneck_tensors(node_dict)

  def _try_loading_model(self, model_path):
    """ Load model from model_path.

    TF models are often saved in one of the three major formats:
      1) Checkpoints with ckpt.meta, ckpt.data, and ckpt.index.
      2) SavedModel format with saved_model.pb and variables/.
      3) Frozen graph in .pb or .pbtxt format.
    When model_path is specified, model is loaded in one of the
    three formats depending on the model_path. When model_path is
    ommitted, child wrapper is responsible for loading the model.
    """
    try:
      self.sess = tf.Session(graph=tf.Graph())
      with self.sess.graph.as_default():
        if tf.io.gfile.isdir(model_path):
          ckpt = tf.train.latest_checkpoint(model_path)
          if ckpt:
            tf.logging.info('Loading from the latest checkpoint.')
            saver = tf.train.import_meta_graph(ckpt + '.meta')
            saver.restore(self.sess, ckpt)
          else:
            tf.logging.info('Loading from SavedModel dir.')
            tf.saved_model.loader.load(self.sess, ['serve'], model_path)
        else:
          input_graph_def = tf.GraphDef()
          if model_path.endswith('.pb'):
            tf.logging.info('Loading from frozen binary graph.')
            with tf.io.gfile.GFile(model_path, 'rb') as f:
              input_graph_def.ParseFromString(f.read())
          else:
            tf.logging.info('Loading from frozen text graph.')
            with tf.io.gfile.GFile(model_path) as f:
              text_format.Parse(f.read(), input_graph_def)
          tf.import_graph_def(input_graph_def)
          self.import_prefix = True

    except Exception as e:
      template = 'An exception of type {0} occurred ' \
                 'when trying to load model from {1}. ' \
                 'Arguments:\n{2!r}'
      tf.logging.warn(template.format(type(e).__name__, model_path, e.args))

  def _find_ends_and_bottleneck_tensors(self, node_dict):
    """ Find tensors from the graph by their names.

    Depending on how the model is loaded, tensors in the graph
    may or may not have 'import/' prefix added to every tensor name.
    This is true even if the tensors already have 'import/' prefix.
    The 'ends' and 'bottlenecks_tensors' dictionary should map to tensors
    with the according name.
    """
    self.bottlenecks_tensors = {}
    self.ends = {}
    for k, v in six.iteritems(node_dict):
      if self.import_prefix:
        v = 'import/' + v
      tensor = self.sess.graph.get_operation_by_name(v.strip(':0')).outputs[0]
      if k == 'input' or k == 'prediction':
        self.ends[k] = tensor
      else:
        self.bottlenecks_tensors[k] = tensor

  def _make_gradient_tensors(self):
    """Makes gradient tensors for all bottleneck tensors.
    """
    self.bottlenecks_gradients = {}
    for bn in self.bottlenecks_tensors:
      self.bottlenecks_gradients[bn] = tf.gradients(
          self.loss, self.bottlenecks_tensors[bn])[0]

  def get_gradient(self, acts, y, bottleneck_name, example):
    """Return the gradient of the loss with respect to the bottleneck_name.

    Args:
      acts: activation of the bottleneck
      y: index of the logit layer
      bottleneck_name: name of the bottleneck to get gradient wrt.
      example: input example. Unused by default. Necessary for getting gradients
        from certain models, such as BERT.

    Returns:
      the gradient array.
    """
    return self.sess.run(self.bottlenecks_gradients[bottleneck_name], {
        self.bottlenecks_tensors[bottleneck_name]: acts,
        self.y_input: y
    })

  def get_predictions(self, examples):
    """Get prediction of the examples.

    Args:
      imgs: array of examples to get predictions

    Returns:
      array of predictions
    """
    return self.adjust_prediction(
        self.sess.run(self.ends['prediction'], {self.ends['input']: examples}))

  def adjust_prediction(self, pred_t):
    """Adjust the prediction tensor to be the expected shape.

    Defaults to a no-op, but necessary to override for GoogleNet
    Returns:
      pred_t: pred_tensor.
    """
    return pred_t

  def reshape_activations(self, layer_acts):
    """Reshapes layer activations as needed to feed through the model network.

    Override this for models that require reshaping of the activations for use
    in TCAV.

    Args:
      layer_acts: Activations as returned by run_examples.

    Returns:
      Activations in model-dependent form; the default is a squeezed array (i.e.
      at most one dimensions of size 1).
    """
    return np.asarray(layer_acts).squeeze()

  def label_to_id(self, label):
    """Convert label (string) to index in the logit layer (id).

    Override this method if label to id mapping is known. Otherwise,
    default id 0 is used.
    """
    tf.logging.warn('label_to_id undefined. Defaults to returning 0.')
    return 0

  def id_to_label(self, idx):
    """Convert index in the logit layer (id) to label (string).

    Override this method if id to label mapping is known.
    """
    return str(idx)

  def run_examples(self, examples, bottleneck_name):
    """Get activations at a bottleneck for provided examples.

    Args:
      examples: example data to feed into network.
      bottleneck_name: string, should be key of self.bottlenecks_tensors

    Returns:
      Activations in the given layer.
    """
    return self.sess.run(self.bottlenecks_tensors[bottleneck_name],
                         {self.ends['input']: examples})

class ImageModelWrapper(ModelWrapper):
  """Wrapper base class for image models."""

  def __init__(self, image_shape):
    super(ModelWrapper, self).__init__()
    # shape of the input image in this model
    self.image_shape = image_shape

  def get_image_shape(self):
    """returns the shape of an input image."""
    return self.image_shape

class PublicImageModelWrapper(ImageModelWrapper):
  """Simple wrapper of the public image models with session object.
  """
  def __init__(self,
               sess,
               model_fn_path,
               labels_path,
               image_shape,
               endpoints_dict,
               scope):
    super(PublicImageModelWrapper, self).__init__(image_shape)
    self.labels = tf.io.gfile.GFile(labels_path).read().splitlines()
    self.ends = PublicImageModelWrapper.import_graph(model_fn_path,
                                                     endpoints_dict,
                                                     self.image_value_range,
                                                     scope=scope)
    self.bottlenecks_tensors = PublicImageModelWrapper.get_bottleneck_tensors(
        scope)
    graph = tf.get_default_graph()

    # Construct gradient ops.
    with graph.as_default():
      self.y_input = tf.placeholder(tf.int64, shape=[None])

      self.pred = tf.expand_dims(self.ends['prediction'][0], 0)
      self.loss = tf.reduce_mean(
          tf.nn.softmax_cross_entropy_with_logits(
              labels=tf.one_hot(
                  self.y_input,
                  self.ends['prediction'].get_shape().as_list()[1]),
              logits=self.pred))
    self._make_gradient_tensors()

  def id_to_label(self, idx):
    return self.labels[idx]

  def label_to_id(self, label):
    return self.labels.index(label)

  @staticmethod
  def create_input(t_input, image_value_range):
    """Create input tensor."""
    def forget_xy(t):
      """Forget sizes of dimensions [1, 2] of a 4d tensor."""
      zero = tf.identity(0)
      return t[:, zero:, zero:, :]

    t_prep_input = t_input
    if len(t_prep_input.shape) == 3:
      t_prep_input = tf.expand_dims(t_prep_input, 0)
    t_prep_input = forget_xy(t_prep_input)
    lo, hi = image_value_range
    t_prep_input = lo + t_prep_input * (hi-lo)
    return t_input, t_prep_input

  # From Alex's code.
  @staticmethod
  def get_bottleneck_tensors(scope):
    """Add Inception bottlenecks and their pre-Relu versions to endpoints dict."""
    graph = tf.get_default_graph()
    bn_endpoints = {}
    for op in graph.get_operations():
      if op.name.startswith(scope+'/') and 'Concat' in op.type:
        name = op.name.split('/')[1]
        bn_endpoints[name] = op.outputs[0]
    return bn_endpoints

  # Load graph and import into graph used by our session
  @staticmethod
  def import_graph(saved_path, endpoints, image_value_range, scope='import'):
    t_input = tf.placeholder(np.float32, [None, None, None, 3])
    graph = tf.Graph()
    assert graph.unique_name(scope, False) == scope, (
        'Scope "%s" already exists. Provide explicit scope names when '
        'importing multiple instances of the model.') % scope

    graph_def = tf.GraphDef.FromString(tf.io.gfile.GFile(saved_path, 'rb').read())

    with tf.name_scope(scope) as sc:
      t_input, t_prep_input = PublicImageModelWrapper.create_input(
          t_input, image_value_range)

      graph_inputs = {}
      graph_inputs[endpoints['input']] = t_prep_input
      myendpoints = tf.import_graph_def(
          graph_def, graph_inputs, list(endpoints.values()), name=sc)
      myendpoints = dict(list(zip(list(endpoints.keys()), myendpoints)))
      myendpoints['input'] = t_input
    return myendpoints

class GoogleNetWrapper_public(PublicImageModelWrapper):

  def __init__(self, sess, model_saved_path, labels_path):
    image_shape_v1 = [224, 224, 3]
    self.image_value_range = (-117, 255-117)
    endpoints_v1 = dict(
        input='input:0',
        logit='softmax2_pre_activation:0',
        prediction='output2:0',
        pre_avgpool='mixed5b:0',
        logit_weight='softmax2_w:0',
        logit_bias='softmax2_b:0',
    )
    self.sess = sess
    super(GoogleNetWrapper_public, self).__init__(sess,
                                                  model_saved_path,
                                                  labels_path,
                                                  image_shape_v1,
                                                  endpoints_v1,
                                                  scope='v1')
    self.model_name = 'GoogleNet_public'

  def adjust_prediction(self, pred_t):
    # Each pred outputs 16, 1008 matrix. The prediction value is the first row.
    # Following tfzoo convention.
    return pred_t[::16]

class InceptionV3Wrapper_public(PublicImageModelWrapper):
  def __init__(self, sess, model_saved_path, labels_path):
    self.image_value_range = (-1, 1)
    image_shape_v3 = [299, 299, 3]
    endpoints_v3 = dict(
        input='Mul:0',
        logit='softmax/logits:0',
        prediction='softmax:0',
        pre_avgpool='mixed_10/join:0',
        logit_weight='softmax/weights:0',
        logit_bias='softmax/biases:0',
    )

    self.sess = sess
    super(InceptionV3Wrapper_public, self).__init__(sess,
                                                  model_saved_path,
                                                  labels_path,
                                                  image_shape_v3,
                                                  endpoints_v3,
                                                  scope='v3')
    self.model_name = 'InceptionV3_public'
<|endoftext|>


<commit_before>import datetime as dt
import logging

from flask import (Blueprint, flash, redirect, render_template,
                   request, session, url_for, jsonify)

from flask_login import login_required, current_user

from server.models import db, Client, Token, Grant
from server.extensions import csrf, oauth_provider

logger = logging.getLogger(__name__)

oauth = Blueprint('oauth', __name__)

@oauth.record
def record_params(setup_state):
    """ Load used app configs into local config on registration from
    server/__init__.py """
    app = setup_state.app
    oauth_provider.init_app(app)

@oauth_provider.clientgetter
def load_client(client_id):
    return Client.query.filter_by(client_id=client_id).first()

@oauth_provider.grantgetter
def load_grant(client_id, code):
    return Grant.query.filter_by(client_id=client_id, code=code).first()

@oauth_provider.grantsetter
def save_grant(client_id, code, request, *args, **kwargs):
    expires = dt.datetime.utcnow() + dt.timedelta(seconds=100)
    grant = Grant(
        client_id=client_id,
        code=code['code'],
        redirect_uri=request.redirect_uri,
        scopes=request.scopes,
        user=current_user,
        expires=expires
    )
    db.session.add(grant)
    db.session.commit()
    return grant

@oauth_provider.tokengetter
def load_token(access_token=None, refresh_token=None):
    if access_token:
        return Token.query.filter_by(access_token=access_token).first()
    elif refresh_token:
        return Token.query.filter_by(refresh_token=refresh_token).first()

@oauth_provider.tokensetter
def save_token(token, orequest, *args, **kwargs):
    toks = Token.query.filter_by(client_id=orequest.client.client_id,
                                 user_id=orequest.user.id).all()
    # make sure that every client has only one token connected to a user
    for t in toks:
        db.session.delete(t)

    expires_in = token.get('expires_in')
    expires = dt.datetime.utcnow() + dt.timedelta(seconds=expires_in)

    tok = Token(
        access_token=token['access_token'],
        refresh_token=token['refresh_token'],
        token_type=token['token_type'],
        scopes=token['scope'].split(),
        expires=expires,
        client_id=orequest.client.client_id,
        user_id=orequest.user.id,
    )
    db.session.add(tok)
    db.session.commit()
    return tok

@oauth.route('/oauth/authorize', methods=['GET', 'POST'])
@oauth_provider.authorize_handler
@login_required
def authorize(*args, **kwargs):
    # Only CSRF protect this route.
    csrf.protect()

    if request.method == 'GET':
        client_id = kwargs.get('client_id')
        client = Client.query.filter_by(client_id=client_id).first()
        kwargs['client'] = client
        return render_template('auth/oauthorize.html', **kwargs)

    confirm = request.form.get('confirm', 'no')
    return confirm == 'yes'

@oauth.route('/oauth/token', methods=['POST'])
@oauth_provider.token_handler
def access_token():
    """ Exchange/Refresh the token. Flask-OAuthLib handles this. """
    return None

@oauth.route('/oauth/revoke', methods=['POST'])
@oauth_provider.revoke_handler
def revoke_token():
    return

@oauth.route('/oauth/errors')
def oauth_errors():
    error = request.args.get('error', 'Unknown Error')
    description = request.args.get('error_description', 'No details available')
    return render_template('errors/generic.html',
                           error=error, description=description), 400
<commit_msg>wild stab at fixing test because csrf on testing is disabled<commit_after>import datetime as dt
import logging

from flask import (Blueprint, flash, redirect, render_template,
                   request, session, url_for, jsonify)

from flask_login import login_required, current_user

from server.models import db, Client, Token, Grant
from server.extensions import csrf, oauth_provider
from server.controllers.auth import use_testing_login

logger = logging.getLogger(__name__)

oauth = Blueprint('oauth', __name__)

@oauth.record
def record_params(setup_state):
    """ Load used app configs into local config on registration from
    server/__init__.py """
    app = setup_state.app
    oauth_provider.init_app(app)

@oauth_provider.clientgetter
def load_client(client_id):
    return Client.query.filter_by(client_id=client_id).first()

@oauth_provider.grantgetter
def load_grant(client_id, code):
    return Grant.query.filter_by(client_id=client_id, code=code).first()

@oauth_provider.grantsetter
def save_grant(client_id, code, request, *args, **kwargs):
    expires = dt.datetime.utcnow() + dt.timedelta(seconds=100)
    grant = Grant(
        client_id=client_id,
        code=code['code'],
        redirect_uri=request.redirect_uri,
        scopes=request.scopes,
        user=current_user,
        expires=expires
    )
    db.session.add(grant)
    db.session.commit()
    return grant

@oauth_provider.tokengetter
def load_token(access_token=None, refresh_token=None):
    if access_token:
        return Token.query.filter_by(access_token=access_token).first()
    elif refresh_token:
        return Token.query.filter_by(refresh_token=refresh_token).first()

@oauth_provider.tokensetter
def save_token(token, orequest, *args, **kwargs):
    toks = Token.query.filter_by(client_id=orequest.client.client_id,
                                 user_id=orequest.user.id).all()
    # make sure that every client has only one token connected to a user
    for t in toks:
        db.session.delete(t)

    expires_in = token.get('expires_in')
    expires = dt.datetime.utcnow() + dt.timedelta(seconds=expires_in)

    tok = Token(
        access_token=token['access_token'],
        refresh_token=token['refresh_token'],
        token_type=token['token_type'],
        scopes=token['scope'].split(),
        expires=expires,
        client_id=orequest.client.client_id,
        user_id=orequest.user.id,
    )
    db.session.add(tok)
    db.session.commit()
    return tok

@oauth.route('/oauth/authorize', methods=['GET', 'POST'])
@oauth_provider.authorize_handler
@login_required
def authorize(*args, **kwargs):
    # Only CSRF protect this route.
    if not use_testing_login():
        csrf.protect()

    if request.method == 'GET':
        client_id = kwargs.get('client_id')
        client = Client.query.filter_by(client_id=client_id).first()
        kwargs['client'] = client
        return render_template('auth/oauthorize.html', **kwargs)

    confirm = request.form.get('confirm', 'no')
    return confirm == 'yes'

@oauth.route('/oauth/token', methods=['POST'])
@oauth_provider.token_handler
def access_token():
    """ Exchange/Refresh the token. Flask-OAuthLib handles this. """
    return None

@oauth.route('/oauth/revoke', methods=['POST'])
@oauth_provider.revoke_handler
def revoke_token():
    return

@oauth.route('/oauth/errors')
def oauth_errors():
    error = request.args.get('error', 'Unknown Error')
    description = request.args.get('error_description', 'No details available')
    return render_template('errors/generic.html',
                           error=error, description=description), 400
<|endoftext|>