<commit_before>import contextlib
import json
import pathlib
import shutil
import subprocess
import time

from attr import attributes, attrib
from hamcrest import assert_that, equal_to, less_than_or_equal_to

import requests

##############################
# General utilities
##############################
_TEST_DIR = pathlib.Path(__file__).parent.parent
_REPO_DIR = _TEST_DIR.parent

# Command execution helper
def _run_command(cmd, work_dir, ignore_errors):
    print("  Running {} in {}".format(cmd, work_dir))
    output = None
    try:
        output = subprocess.check_output(
            cmd, cwd=work_dir, stderr=subprocess.PIPE
        ).decode()
    except subprocess.CalledProcessError as exc:
        output = exc.output.decode()
        if not ignore_errors:
            print("=== stdout for failed command ===")
            print(output)
            print("=== stderr for failed command ===")
            print(exc.stderr.decode())
            raise
    return output

##############################
# Local VM management
##############################

_VM_HOSTNAME_PREFIX = "leapp-tests-"
_VM_DEFS = {
    _VM_HOSTNAME_PREFIX + path.name: str(path)
        for path in (_TEST_DIR / "vmdefs").iterdir()
}

class VirtualMachineHelper(object):
    """Test step helper to launch and manage VMs

    Currently based specifically on local Vagrant VMs
    """

    def __init__(self):
        self.machines = {}
        self._resource_manager = contextlib.ExitStack()

    def ensure_local_vm(self, name, definition, destroy=False):
        """Ensure a local VM exists based on the given definition

        *name*: name used to refer to the VM in scenario steps
        *definition*: directory name in integration-tests/vmdefs
        *destroy*: whether or not to destroy any existing VM
        """
        hostname = _VM_HOSTNAME_PREFIX + definition
        if hostname not in _VM_DEFS:
            raise ValueError("Unknown VM image: {}".format(definition))
        if destroy:
            self._vm_destroy(hostname)
        self._vm_up(name, hostname)
        if destroy:
            self._resource_manager.callback(self._vm_destroy, name)
        else:
            self._resource_manager.callback(self._vm_halt, name)

    def get_hostname(self, name):
        """Return the expected hostname for the named machine"""
        return self.machines[name]

    def close(self):
        """Halt or destroy all created VMs"""
        self._resource_manager.close()

    @staticmethod
    def _run_vagrant(hostname, *args, ignore_errors=False):
        # TODO: explore https://pypi.python.org/pypi/python-vagrant
        vm_dir = _VM_DEFS[hostname]
        cmd = ["vagrant"]
        cmd.extend(args)
        return _run_command(cmd, vm_dir, ignore_errors)

    def _vm_up(self, name, hostname):
        result = self._run_vagrant(hostname, "up", "--provision")
        print("Started {} VM instance".format(hostname))
        self.machines[name] = hostname
        return result

    def _vm_halt(self, name):
        hostname = self.machines.pop(name)
        result = self._run_vagrant(hostname, "halt", ignore_errors=True)
        print("Suspended {} VM instance".format(hostname))
        return result

    def _vm_destroy(self, name):
        hostname = self.machines.pop(name)
        result = self._run_vagrant(hostname, "destroy", ignore_errors=True)
        print("Destroyed {} VM instance".format(hostname))
        return result

##############################
# Leapp commands
##############################

_LEAPP_SRC_DIR = _REPO_DIR / "src"
_LEAPP_BIN_DIR = _REPO_DIR / "bin"
_LEAPP_TOOL = str(_LEAPP_BIN_DIR / "leapp-tool")

_SSH_IDENTITY = str(_REPO_DIR / "integration-tests/config/leappto_testing_key")

def _install_client():
    """Install the CLI and its dependencies into a Python 2.7 environment"""
    py27 = shutil.which("python2.7")
    base_cmd = ["pipsi", "--bin-dir", str(_LEAPP_BIN_DIR)]
    if pathlib.Path(_LEAPP_TOOL).exists():
        # For some reason, `upgrade` returns 1 even though it appears to work
        # so we instead do a full uninstall/reinstall before the test run
        uninstall = base_cmd + ["uninstall", "--yes", "leappto"]
        _run_command(uninstall, work_dir=str(_REPO_DIR), ignore_errors=False)
    install = base_cmd + ["install", "--python", py27, str(_LEAPP_SRC_DIR)]
    print(_run_command(install, work_dir=str(_REPO_DIR), ignore_errors=False))

@attributes
class MigrationInfo(object):
    """Details of local hosts involved in an app migration command

    *local_vm_count*: Total number of local VMs found during migration
    *source_ip*: host accessible IP address found for source VM
    *target_ip*: host accessible IP address found for target VM
    """
    local_vm_count = attrib()
    source_ip = attrib()
    target_ip = attrib()

    @classmethod
    def from_vm_list(cls, machines, source_host, target_host):
        """Build a result given a local VM listing and migration hostnames"""
        vm_count = len(machines)
        source_ip = target_ip = None
        for machine in machines:
            if machine["hostname"] == source_host:
                source_ip = machine["ip"][0]
            if machine["hostname"] == target_host:
                target_ip = machine["ip"][0]
            if source_ip is not None and target_ip is not None:
                break
        return cls(vm_count, source_ip, target_ip)

class ClientHelper(object):
    """Test step helper to invoke the LeApp CLI

    Requires a VirtualMachineHelper instance
    """

    def __init__(self, vm_helper):
        self._vm_helper = vm_helper

    def redeploy_as_macrocontainer(self, source_vm, target_vm):
        """Recreate source VM as a macrocontainer on given target VM"""
        vm_helper = self._vm_helper
        source_host = vm_helper.get_hostname(source_vm)
        target_host = vm_helper.get_hostname(target_vm)
        self._convert_vm_to_macrocontainer(source_host, target_host)
        return self._get_migration_host_info(source_host, target_host)

    def check_response_time(self, cmd_args, time_limit, complete_identity=False):
        """Check given command completes within the specified time limit

        Returns the contents of stdout as a string.
        """
        if complete_identity:
            cmd_args.extend(['--identity', _SSH_IDENTITY])
        start = time.monotonic()
        cmd_output = self._run_leapp(*cmd_args)
        response_time = time.monotonic() - start
        assert_that(response_time, less_than_or_equal_to(time_limit))
        return cmd_output

    @staticmethod
    def _run_leapp(*args):
        cmd = [_LEAPP_TOOL]
        cmd.extend(args)
        return _run_command(cmd, work_dir=str(_LEAPP_BIN_DIR), ignore_errors=False)

    @classmethod
    def _convert_vm_to_macrocontainer(cls, source_host, target_host):
        result = cls._run_leapp(
            "migrate-machine",
            "--identity", _SSH_IDENTITY,
            "-t", target_host, source_host
        )
        msg = "Redeployed {} as macrocontainer on {}"
        print(msg.format(source_host, target_host))
        return result

    @classmethod
    def _get_migration_host_info(cls, source_host, target_host):
        leapp_output = cls._run_leapp("list-machines", "--shallow")
        machines = json.loads(leapp_output)["machines"]
        return MigrationInfo.from_vm_list(machines, source_host, target_host)

##############################
# Service status checking
##############################

class RequestsHelper(object):
    """Test step helper to check HTTP responses"""

    @classmethod
    def get_response(cls, service_url, wait_for_connection=None):
        """Get HTTP response from given service URL

        Responses are returned as requests.Response objects

        *service_url*: the service URL to query
        *wait_for_connection*: number of seconds to wait for a HTTP connection
                               to the service. `None` indicates that a response
                               is expected immediately.
        """
        deadline = time.monotonic()
        if wait_for_connection is None:
            fail_msg = "No response from {}".format(service_url)
        else:
            fail_msg = "No response from {} within {} seconds".format(
                service_url,
                wait_for_connection
            )
            deadline += wait_for_connection
        while True:
            try:
                return requests.get(service_url)
            except Exception:
                pass
            if time.monotonic() >= deadline:
                break
        raise AssertionError(fail_msg)

    @classmethod
    def get_responses(cls, urls_to_check):
        """Check responses from multiple given URLs

        Each URL can be either a string (which will be expected to return
        a response immediately), or else a (service_url, wait_for_connection)
        pair, which is interpreted as described for `get_response()`.

        Response are returned as a dictionary mapping from the service URLs
        to requests.Response objects.
        """
        # TODO: Use concurrent.futures to check the given URLs in parallel
        responses = {}
        for url_to_check in urls_to_check:
            if isinstance(url_to_check, tuple):
                url_to_check, wait_for_connection = url_to_check
            else:
                wait_for_connection = None
            responses[url_to_check] = cls.get_response(url_to_check,
                                                       wait_for_connection)
        return responses

    @classmethod
    def compare_redeployed_response(cls, original_ip, redeployed_ip, *,
                                    tcp_port, status, wait_for_target):
        """Compare a pre-migration app response with a redeployed response

        Expects an immediate response from the original IP, and allows for
        a delay before the redeployment target starts returning responses
        """
        # Get response from source VM
        original_url = "http://{}:{}".format(original_ip, tcp_port)
        original_response = cls.get_response(original_url)
        print("Response received from {}".format(original_url))
        original_status = original_response.status_code
        assert_that(original_status, equal_to(status), "Original status")
        # Get response from target VM
        redeployed_url = "http://{}:{}".format(redeployed_ip, tcp_port)
        redeployed_response = cls.get_response(redeployed_url, wait_for_target)
        print("Response received from {}".format(redeployed_url))
        # Compare the responses
        assert_that(redeployed_response.status_code, equal_to(original_status), "Redeployed status")
        original_data = original_response.text
        redeployed_data = redeployed_response.text
        assert_that(redeployed_data, equal_to(original_data), "Same response")

##############################
# Test execution hooks
##############################

# The @skip support here is based on:
# http://stackoverflow.com/questions/36482419/how-do-i-skip-a-test-in-the-behave-python-bdd-framework/42721605#42721605
_AUTOSKIP_TAG = "skip"
_WIP_TAG = "wip"
def _skip_test_group(context, test_group):
    """Decides whether or not to skip a test feature or test scenario

    Test groups are skipped if they're tagged with `@skip` and the `skip` tag
    is not explicitly requested through the command line options.

    The `--wip` option currently overrides `--tags`, so groups tagged with both
    `@skip` *and* `@wip` are also executed when the `wip` tag is requested.
    """
    active_tags = context.config.tags
    autoskip_tag_set = _AUTOSKIP_TAG in test_group.tags
    autoskip_tag_requested = active_tags and active_tags.check([_AUTOSKIP_TAG])
    wip_tag_set = _WIP_TAG in test_group.tags
    wip_tag_requested = active_tags and active_tags.check([_AUTOSKIP_TAG, _WIP_TAG])
    override_autoskip = autoskip_tag_requested or (wip_tag_set and wip_tag_requested)
    skip_group = autoskip_tag_set and not override_autoskip
    if skip_group:
        test_group.skip("Marked with @skip")
    return skip_group

def before_all(context):
    # Basic info about the test repository
    context.BASE_REPO_DIR = _REPO_DIR
    context.BASE_TEST_DIR = _TEST_DIR

    # Some steps require sudo, so for convenience in interactive use,
    # we ensure we prompt for elevated permissions immediately,
    # rather than potentially halting midway through a test
    subprocess.check_output(["sudo", "echo", "Elevated permissions needed"])

    # Install the CLI for use in the tests
    _install_client()

    # Use contextlib.ExitStack to manage global resources
    context._global_cleanup = contextlib.ExitStack()

def before_feature(context, feature):
    if _skip_test_group(context, feature):
        return

    # Use contextlib.ExitStack to manage per-feature resources
    context._feature_cleanup = contextlib.ExitStack()

def before_scenario(context, scenario):
    if _skip_test_group(context, scenario):
        return

    # Each scenario has a contextlib.ExitStack instance for resource cleanup
    context.scenario_cleanup = contextlib.ExitStack()

    # Each scenario gets a VirtualMachineHelper instance
    # VMs are relatively slow to start/stop, so by default, we defer halting
    # or destroying them to the end of the overall test run
    # Feature steps can still opt in to eagerly cleaning up a scenario's VMs
    # by doing `context.scenario_cleanup.callback(context.vm_helper.close)`
    context.vm_helper = vm_helper = VirtualMachineHelper()
    context._global_cleanup.callback(vm_helper.close)

    # Each scenario gets a ClientHelper instance
    context.cli_helper = cli_helper = ClientHelper(context.vm_helper)

    # Each scenario gets a RequestsHelper instance
    context.http_helper = RequestsHelper()

def after_scenario(context, scenario):
    if scenario.status == "skipped":
        return
    context.scenario_cleanup.close()

def after_feature(context, feature):
    if feature.status == "skipped":
        return
    context._feature_cleanup.close()

def after_all(context):
    context._global_cleanup.close()
<commit_msg>Fix permissions on testing key prior to run<commit_after>import contextlib
import json
import os
import pathlib
import shutil
import subprocess
import time

from attr import attributes, attrib
from hamcrest import assert_that, equal_to, less_than_or_equal_to

import requests

##############################
# General utilities
##############################
_TEST_DIR = pathlib.Path(__file__).parent.parent
_REPO_DIR = _TEST_DIR.parent

# Command execution helper
def _run_command(cmd, work_dir, ignore_errors):
    print("  Running {} in {}".format(cmd, work_dir))
    output = None
    try:
        output = subprocess.check_output(
            cmd, cwd=work_dir, stderr=subprocess.PIPE
        ).decode()
    except subprocess.CalledProcessError as exc:
        output = exc.output.decode()
        if not ignore_errors:
            print("=== stdout for failed command ===")
            print(output)
            print("=== stderr for failed command ===")
            print(exc.stderr.decode())
            raise
    return output

##############################
# Local VM management
##############################

_VM_HOSTNAME_PREFIX = "leapp-tests-"
_VM_DEFS = {
    _VM_HOSTNAME_PREFIX + path.name: str(path)
        for path in (_TEST_DIR / "vmdefs").iterdir()
}

class VirtualMachineHelper(object):
    """Test step helper to launch and manage VMs

    Currently based specifically on local Vagrant VMs
    """

    def __init__(self):
        self.machines = {}
        self._resource_manager = contextlib.ExitStack()

    def ensure_local_vm(self, name, definition, destroy=False):
        """Ensure a local VM exists based on the given definition

        *name*: name used to refer to the VM in scenario steps
        *definition*: directory name in integration-tests/vmdefs
        *destroy*: whether or not to destroy any existing VM
        """
        hostname = _VM_HOSTNAME_PREFIX + definition
        if hostname not in _VM_DEFS:
            raise ValueError("Unknown VM image: {}".format(definition))
        if destroy:
            self._vm_destroy(hostname)
        self._vm_up(name, hostname)
        if destroy:
            self._resource_manager.callback(self._vm_destroy, name)
        else:
            self._resource_manager.callback(self._vm_halt, name)

    def get_hostname(self, name):
        """Return the expected hostname for the named machine"""
        return self.machines[name]

    def close(self):
        """Halt or destroy all created VMs"""
        self._resource_manager.close()

    @staticmethod
    def _run_vagrant(hostname, *args, ignore_errors=False):
        # TODO: explore https://pypi.python.org/pypi/python-vagrant
        vm_dir = _VM_DEFS[hostname]
        cmd = ["vagrant"]
        cmd.extend(args)
        return _run_command(cmd, vm_dir, ignore_errors)

    def _vm_up(self, name, hostname):
        result = self._run_vagrant(hostname, "up", "--provision")
        print("Started {} VM instance".format(hostname))
        self.machines[name] = hostname
        return result

    def _vm_halt(self, name):
        hostname = self.machines.pop(name)
        result = self._run_vagrant(hostname, "halt", ignore_errors=True)
        print("Suspended {} VM instance".format(hostname))
        return result

    def _vm_destroy(self, name):
        hostname = self.machines.pop(name)
        result = self._run_vagrant(hostname, "destroy", ignore_errors=True)
        print("Destroyed {} VM instance".format(hostname))
        return result

##############################
# Leapp commands
##############################

_LEAPP_SRC_DIR = _REPO_DIR / "src"
_LEAPP_BIN_DIR = _REPO_DIR / "bin"
_LEAPP_TOOL = str(_LEAPP_BIN_DIR / "leapp-tool")

_SSH_IDENTITY = str(_REPO_DIR / "integration-tests/config/leappto_testing_key")

def _install_client():
    """Install the CLI and its dependencies into a Python 2.7 environment"""
    py27 = shutil.which("python2.7")
    base_cmd = ["pipsi", "--bin-dir", str(_LEAPP_BIN_DIR)]
    if pathlib.Path(_LEAPP_TOOL).exists():
        # For some reason, `upgrade` returns 1 even though it appears to work
        # so we instead do a full uninstall/reinstall before the test run
        uninstall = base_cmd + ["uninstall", "--yes", "leappto"]
        _run_command(uninstall, work_dir=str(_REPO_DIR), ignore_errors=False)
    install = base_cmd + ["install", "--python", py27, str(_LEAPP_SRC_DIR)]
    print(_run_command(install, work_dir=str(_REPO_DIR), ignore_errors=False))
    # Ensure private SSH key is only readable by the current user
    os.chmod(_SSH_IDENTITY, 0o600)

@attributes
class MigrationInfo(object):
    """Details of local hosts involved in an app migration command

    *local_vm_count*: Total number of local VMs found during migration
    *source_ip*: host accessible IP address found for source VM
    *target_ip*: host accessible IP address found for target VM
    """
    local_vm_count = attrib()
    source_ip = attrib()
    target_ip = attrib()

    @classmethod
    def from_vm_list(cls, machines, source_host, target_host):
        """Build a result given a local VM listing and migration hostnames"""
        vm_count = len(machines)
        source_ip = target_ip = None
        for machine in machines:
            if machine["hostname"] == source_host:
                source_ip = machine["ip"][0]
            if machine["hostname"] == target_host:
                target_ip = machine["ip"][0]
            if source_ip is not None and target_ip is not None:
                break
        return cls(vm_count, source_ip, target_ip)

class ClientHelper(object):
    """Test step helper to invoke the LeApp CLI

    Requires a VirtualMachineHelper instance
    """

    def __init__(self, vm_helper):
        self._vm_helper = vm_helper

    def redeploy_as_macrocontainer(self, source_vm, target_vm):
        """Recreate source VM as a macrocontainer on given target VM"""
        vm_helper = self._vm_helper
        source_host = vm_helper.get_hostname(source_vm)
        target_host = vm_helper.get_hostname(target_vm)
        self._convert_vm_to_macrocontainer(source_host, target_host)
        return self._get_migration_host_info(source_host, target_host)

    def check_response_time(self, cmd_args, time_limit, complete_identity=False):
        """Check given command completes within the specified time limit

        Returns the contents of stdout as a string.
        """
        if complete_identity:
            cmd_args.extend(['--identity', _SSH_IDENTITY])
        start = time.monotonic()
        cmd_output = self._run_leapp(*cmd_args)
        response_time = time.monotonic() - start
        assert_that(response_time, less_than_or_equal_to(time_limit))
        return cmd_output

    @staticmethod
    def _run_leapp(*args):
        cmd = [_LEAPP_TOOL]
        cmd.extend(args)
        return _run_command(cmd, work_dir=str(_LEAPP_BIN_DIR), ignore_errors=False)

    @classmethod
    def _convert_vm_to_macrocontainer(cls, source_host, target_host):
        result = cls._run_leapp(
            "migrate-machine",
            "--identity", _SSH_IDENTITY,
            "-t", target_host, source_host
        )
        msg = "Redeployed {} as macrocontainer on {}"
        print(msg.format(source_host, target_host))
        return result

    @classmethod
    def _get_migration_host_info(cls, source_host, target_host):
        leapp_output = cls._run_leapp("list-machines", "--shallow")
        machines = json.loads(leapp_output)["machines"]
        return MigrationInfo.from_vm_list(machines, source_host, target_host)

##############################
# Service status checking
##############################

class RequestsHelper(object):
    """Test step helper to check HTTP responses"""

    @classmethod
    def get_response(cls, service_url, wait_for_connection=None):
        """Get HTTP response from given service URL

        Responses are returned as requests.Response objects

        *service_url*: the service URL to query
        *wait_for_connection*: number of seconds to wait for a HTTP connection
                               to the service. `None` indicates that a response
                               is expected immediately.
        """
        deadline = time.monotonic()
        if wait_for_connection is None:
            fail_msg = "No response from {}".format(service_url)
        else:
            fail_msg = "No response from {} within {} seconds".format(
                service_url,
                wait_for_connection
            )
            deadline += wait_for_connection
        while True:
            try:
                return requests.get(service_url)
            except Exception:
                pass
            if time.monotonic() >= deadline:
                break
        raise AssertionError(fail_msg)

    @classmethod
    def get_responses(cls, urls_to_check):
        """Check responses from multiple given URLs

        Each URL can be either a string (which will be expected to return
        a response immediately), or else a (service_url, wait_for_connection)
        pair, which is interpreted as described for `get_response()`.

        Response are returned as a dictionary mapping from the service URLs
        to requests.Response objects.
        """
        # TODO: Use concurrent.futures to check the given URLs in parallel
        responses = {}
        for url_to_check in urls_to_check:
            if isinstance(url_to_check, tuple):
                url_to_check, wait_for_connection = url_to_check
            else:
                wait_for_connection = None
            responses[url_to_check] = cls.get_response(url_to_check,
                                                       wait_for_connection)
        return responses

    @classmethod
    def compare_redeployed_response(cls, original_ip, redeployed_ip, *,
                                    tcp_port, status, wait_for_target):
        """Compare a pre-migration app response with a redeployed response

        Expects an immediate response from the original IP, and allows for
        a delay before the redeployment target starts returning responses
        """
        # Get response from source VM
        original_url = "http://{}:{}".format(original_ip, tcp_port)
        original_response = cls.get_response(original_url)
        print("Response received from {}".format(original_url))
        original_status = original_response.status_code
        assert_that(original_status, equal_to(status), "Original status")
        # Get response from target VM
        redeployed_url = "http://{}:{}".format(redeployed_ip, tcp_port)
        redeployed_response = cls.get_response(redeployed_url, wait_for_target)
        print("Response received from {}".format(redeployed_url))
        # Compare the responses
        assert_that(redeployed_response.status_code, equal_to(original_status), "Redeployed status")
        original_data = original_response.text
        redeployed_data = redeployed_response.text
        assert_that(redeployed_data, equal_to(original_data), "Same response")

##############################
# Test execution hooks
##############################

# The @skip support here is based on:
# http://stackoverflow.com/questions/36482419/how-do-i-skip-a-test-in-the-behave-python-bdd-framework/42721605#42721605
_AUTOSKIP_TAG = "skip"
_WIP_TAG = "wip"
def _skip_test_group(context, test_group):
    """Decides whether or not to skip a test feature or test scenario

    Test groups are skipped if they're tagged with `@skip` and the `skip` tag
    is not explicitly requested through the command line options.

    The `--wip` option currently overrides `--tags`, so groups tagged with both
    `@skip` *and* `@wip` are also executed when the `wip` tag is requested.
    """
    active_tags = context.config.tags
    autoskip_tag_set = _AUTOSKIP_TAG in test_group.tags
    autoskip_tag_requested = active_tags and active_tags.check([_AUTOSKIP_TAG])
    wip_tag_set = _WIP_TAG in test_group.tags
    wip_tag_requested = active_tags and active_tags.check([_AUTOSKIP_TAG, _WIP_TAG])
    override_autoskip = autoskip_tag_requested or (wip_tag_set and wip_tag_requested)
    skip_group = autoskip_tag_set and not override_autoskip
    if skip_group:
        test_group.skip("Marked with @skip")
    return skip_group

def before_all(context):
    # Basic info about the test repository
    context.BASE_REPO_DIR = _REPO_DIR
    context.BASE_TEST_DIR = _TEST_DIR

    # Some steps require sudo, so for convenience in interactive use,
    # we ensure we prompt for elevated permissions immediately,
    # rather than potentially halting midway through a test
    subprocess.check_output(["sudo", "echo", "Elevated permissions needed"])

    # Install the CLI for use in the tests
    _install_client()

    # Use contextlib.ExitStack to manage global resources
    context._global_cleanup = contextlib.ExitStack()

def before_feature(context, feature):
    if _skip_test_group(context, feature):
        return

    # Use contextlib.ExitStack to manage per-feature resources
    context._feature_cleanup = contextlib.ExitStack()

def before_scenario(context, scenario):
    if _skip_test_group(context, scenario):
        return

    # Each scenario has a contextlib.ExitStack instance for resource cleanup
    context.scenario_cleanup = contextlib.ExitStack()

    # Each scenario gets a VirtualMachineHelper instance
    # VMs are relatively slow to start/stop, so by default, we defer halting
    # or destroying them to the end of the overall test run
    # Feature steps can still opt in to eagerly cleaning up a scenario's VMs
    # by doing `context.scenario_cleanup.callback(context.vm_helper.close)`
    context.vm_helper = vm_helper = VirtualMachineHelper()
    context._global_cleanup.callback(vm_helper.close)

    # Each scenario gets a ClientHelper instance
    context.cli_helper = cli_helper = ClientHelper(context.vm_helper)

    # Each scenario gets a RequestsHelper instance
    context.http_helper = RequestsHelper()

def after_scenario(context, scenario):
    if scenario.status == "skipped":
        return
    context.scenario_cleanup.close()

def after_feature(context, feature):
    if feature.status == "skipped":
        return
    context._feature_cleanup.close()

def after_all(context):
    context._global_cleanup.close()
<|endoftext|>


<commit_before>import os.path

__all__ = ('__version__', '__build__')

__version__ = (0, 1)

def _get_git_revision(path):
    revision_file = os.path.join(path, 'refs', 'heads', 'master')
    if not os.path.exists(revision_file):
        return None
    fh = open(revision_file, 'r')
    try:
        return fh.read()
    finally:
        fh.close()

def get_revision():
    """
    :returns: Revision number of this branch/checkout, if available. None if
        no revision number can be determined.
    """
    package_dir = os.path.dirname(__file__)
    checkout_dir = os.path.normpath(os.path.join(package_dir, '..'))
    path = os.path.join(checkout_dir, '.git')
    if os.path.exists(path):
        return _get_git_revision(path)
    return None

__build__ = get_revision()<commit_msg>Bump version for PyPi<commit_after>import os.path

__all__ = ('__version__', '__build__')

__version__ = (0, 2)

def _get_git_revision(path):
    revision_file = os.path.join(path, 'refs', 'heads', 'master')
    if not os.path.exists(revision_file):
        return None
    fh = open(revision_file, 'r')
    try:
        return fh.read()
    finally:
        fh.close()

def get_revision():
    """
    :returns: Revision number of this branch/checkout, if available. None if
        no revision number can be determined.
    """
    package_dir = os.path.dirname(__file__)
    checkout_dir = os.path.normpath(os.path.join(package_dir, '..'))
    path = os.path.join(checkout_dir, '.git')
    if os.path.exists(path):
        return _get_git_revision(path)
    return None

__build__ = get_revision()<|endoftext|>


<commit_before>"""APIRequestHandler subclasses for API endpoints."""

from bs4 import BeautifulSoup
from tornado.web import HTTPError, asynchronous
from tornado import gen
import pbkdf2
import logging
import yaml
import urlparse

from feedreader.api_request_handler import APIRequestHandler
from feedreader.database import Feed, User

logger = logging.getLogger(__name__)

class MainHandler(APIRequestHandler):

    def get(self):
        """Return a hello world message."""
        with self.get_db_session() as session:
            username = self.require_auth(session)
            self.write({"message": "Hello, {}.".format(username)})

class UsersHandler(APIRequestHandler):

    def get(self):
        """Return information about the current user."""
        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            username = user.username
        self.write({'username': username})
        self.set_status(200)

    def post(self):
        """Create a new user."""
        body = self.require_body_schema({
            "type": "object",
            "properties": {
                "username": {
                    "type": "string",
                    "pattern": "^[^:\s]*$",
                    "maxLength": 32,
                    "minLength": 1
                },
                "password": {
                    "type": "string",
                    "minLength": 1
                },
            },
            "required": ["username", "password"],
        })

        with self.get_db_session() as session:
            # check if username already exists
            if session.query(User).get(body["username"]) is not None:
                raise HTTPError(400, reason="Username already registered")
            # save new user
            password_hash = pbkdf2.crypt(body["password"])
            new_user = User(body["username"], password_hash)
            session.add(new_user)
            session.commit()
        self.set_status(201)

class FeedsHandler(APIRequestHandler):

    def get(self):
        """Return a list of the user's subscribed feeds."""
        feeds = []
        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            for feed in user.subscriptions:
                feeds.append({
                    'id': feed.id,
                    'name': feed.title,
                    'url': feed.site_url,
                    'image_url': feed.image_url,
                    'unreads': user.get_num_unread_entries(feed),
                })
        self.write({'feeds': feeds})
        self.set_status(200)

    @classmethod
    @gen.coroutine
    def subscribe_feed(cls, session, user, celery_poller, tasks, url):
        """Subscribe the user to a feed and return the feed ID.

        Raises HTTPError if the feed at the given url can't be subscribed to.
        """
        # add a new feed if it doesn't exist already
        feed = session.query(Feed).filter(Feed.feed_url == url).first()
        if feed is None:
            res = yield celery_poller.run_task(tasks.fetch_feed, url)
            res = yaml.safe_load(res)
            if "error" in res:
                logger.warning("Failed to fetch new feed: '{}'"
                               .format(res['error']))
                raise HTTPError(400, reason=res['error'])
            feed = res['feed']
            entries = res['entries']
            parsed_url = urlparse.urlparse(url)
            new_url = urlparse.urlunparse([
                parsed_url.scheme,
                parsed_url.netloc,
                parsed_url.path,
                '',
                '',
                ''
            ])
            new_feed = session.query(Feed).filter(
                Feed.feed_url == new_url
            ).first()
            if new_feed is not None:
                url = new_url
            if new_url != url:
                new_res = yield celery_poller.run_task(
                    tasks.fetch_feed,
                    new_url,
                )
                new_res = yaml.safe_load(new_res)
                if "error" in new_res:
                    logger.warning("Failed to fetch new feed: '{}'"
                               .format(new_res['error']))
                    raise HTTPError(400, reason=new_res['error'])
                new_entries = new_res['entries']
                if len(entries) == len(new_entries):
                    # assume they are the same until proven wrong
                    same = True
                    for i in range(len(entries)):
                        if entries[i].title != new_entries[i].title or\
                           entries[i].content != new_entries[i].content or\
                           entries[i].author != new_entries[i].author or\
                           entries[i].guid != new_entries[i].guid:
                            same = False
                            break
                    if same:
                        feed.feed_url = new_url
            # feed url may have resolved to a feed that already exists
            existing_feed = session.query(Feed)\
                    .filter(Feed.feed_url == feed.feed_url).all()
            if len(existing_feed) > 0:
                feed = existing_feed[0]
            else:
                # only add the first entry with a given guid
                # this assumes the first one is the newest
                entries_by_guid = {}
                for entry in res['entries']:
                    if entry.guid not in entries_by_guid:
                        entries_by_guid[entry.guid] = entry
                feed.add_all(entries_by_guid.values())
                session.commit()
        if user.has_subscription(feed):
            raise HTTPError(400, reason="Already subscribed to feed")

        # subscribe the user to the feed
        user.subscribe(feed)
        session.commit()
        raise gen.Return(feed.id)

    @asynchronous
    @gen.coroutine
    def post(self):
        """Subscribe the user to a new feed."""
        body = self.require_body_schema({
            'type': 'object',
            'properties': {
                'url': {'type': 'string'},
            },
            'required': ['url'],
        })
        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            feed_id = yield self.subscribe_feed(session, user,
                                                self.celery_poller, self.tasks,
                                                body['url'])
            # TODO: technically this is supposed to be an absolute URL
            self.set_header("Location", "/feeds/{}".format(feed_id))
        self.set_status(201)

class FeedHandler(APIRequestHandler):

    def get(self, feed_id):
        """Return metadata for a subscribed feed."""
        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            feed = session.query(Feed).get(int(feed_id))

            # Make sure the feed exists
            if feed is None:
                raise HTTPError(404, reason='This feed does not exist')

            # Make sure the user is subscribed to this feed
            if not user.has_subscription(feed):
                raise HTTPError(404, reason='This feed does not exist')

            self.write({
                'id': feed.id,
                'name': feed.title,
                'url': feed.site_url,
                'image_url': feed.image_url,
                'unreads': user.get_num_unread_entries(feed),
            })
        self.set_status(200)

    def delete(self, feed_id):
        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            feed = session.query(Feed).get(int(feed_id))

            # Make sure the feed exists
            if feed is None:
                raise HTTPError(404, reason='This feed does not exist')

            # Make sure the user is subscribed to this feed
            if not user.has_subscription(feed):
                raise HTTPError(404, reason='This feed does not exist')

            user.unsubscribe(feed)
        self.set_status(204)

class FeedEntriesHandler(APIRequestHandler):

    def get(self, feed_id):
        """Return a list of entry IDs for a feed."""
        entry_filter = self.get_argument('filter', None)

        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            feed = session.query(Feed).get(int(feed_id))

            # Make sure the feed exists
            if feed is None:
                raise HTTPError(404, reason='This feed does not exist')

            # Make sure the user is subscribed to this feed
            if not user.has_subscription(feed):
                raise HTTPError(404, reason='This feed does not exist')

            # Make sure the filter keyword is valid
            if entry_filter not in ['read', 'unread', None]:
                raise HTTPError(400, reason='Filter keyboard is not valid')

            # Get feed entries
            entries = feed.get_entries(user, entry_filter)
            self.write({'entries': [entry.id for entry in entries]})
            self.set_status(200)

class EntriesHandler(APIRequestHandler):

    @classmethod
    @gen.coroutine
    def truncate(cls, content, length):
        return BeautifulSoup(content).get_text()[:length]

    @asynchronous
    @gen.coroutine
    def get(self, entry_ids):
        entry_ids = [int(entry_id) for entry_id in entry_ids.split(",")]

        try:
            length = int(self.get_argument('truncate', 0))
        except ValueError:
            raise HTTPError(400, reason="Invalid truncation size.")

        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            entries = user.get_entries(entry_ids)

            if len(entries) != len(entry_ids):
                raise HTTPError(404, "Entry does not exist.")

            entries_json = []
            for entry in entries:
                entry_json = {
                    "id": entry.id,
                    "title": entry.title,
                    "pub-date": entry.date,
                    "read": user.has_read(entry),
                    "author": entry.author,
                    "feed_id": entry.feed_id,
                    "url": entry.url,
                }

                if length:
                    entry_json['content'] = yield self.truncate(entry.content,
                                                                length)
                else:
                    entry_json['content'] = entry.content
                entries_json.append(entry_json)

            self.write({'entries': entries_json})
        self.set_status(200)

    def patch(self, entry_ids):
        entry_ids = [int(entry_id) for entry_id in entry_ids.split(",")]
        body = self.require_body_schema({
            "type": "object",
            "properties": {
                "read": {"type": "boolean"},
            },
            "required": ["read"],
        })

        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            entries = user.get_entries(entry_ids)

            if len(entries) != len(entry_ids):
                raise HTTPError(404, "Entry does not exist.")

            if body["read"]:
                for entry in entries:
                    if not user.has_read(entry):
                        user.read(entry)
            else:
                for entry in entries:
                    if user.has_read(entry):
                        user.unread(entry)
        self.set_status(200)
<commit_msg>Revert commit that breaks dup detection<commit_after>"""APIRequestHandler subclasses for API endpoints."""

from bs4 import BeautifulSoup
from tornado.web import HTTPError, asynchronous
from tornado import gen
import pbkdf2
import logging
import yaml

from feedreader.api_request_handler import APIRequestHandler
from feedreader.database import Feed, User

logger = logging.getLogger(__name__)

class MainHandler(APIRequestHandler):

    def get(self):
        """Return a hello world message."""
        with self.get_db_session() as session:
            username = self.require_auth(session)
            self.write({"message": "Hello, {}.".format(username)})

class UsersHandler(APIRequestHandler):

    def get(self):
        """Return information about the current user."""
        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            username = user.username
        self.write({'username': username})
        self.set_status(200)

    def post(self):
        """Create a new user."""
        body = self.require_body_schema({
            "type": "object",
            "properties": {
                "username": {
                    "type": "string",
                    "pattern": "^[^:\s]*$",
                    "maxLength": 32,
                    "minLength": 1
                },
                "password": {
                    "type": "string",
                    "minLength": 1
                },
            },
            "required": ["username", "password"],
        })

        with self.get_db_session() as session:
            # check if username already exists
            if session.query(User).get(body["username"]) is not None:
                raise HTTPError(400, reason="Username already registered")
            # save new user
            password_hash = pbkdf2.crypt(body["password"])
            new_user = User(body["username"], password_hash)
            session.add(new_user)
            session.commit()
        self.set_status(201)

class FeedsHandler(APIRequestHandler):

    def get(self):
        """Return a list of the user's subscribed feeds."""
        feeds = []
        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            for feed in user.subscriptions:
                feeds.append({
                    'id': feed.id,
                    'name': feed.title,
                    'url': feed.site_url,
                    'image_url': feed.image_url,
                    'unreads': user.get_num_unread_entries(feed),
                })
        self.write({'feeds': feeds})
        self.set_status(200)

    @classmethod
    @gen.coroutine
    def subscribe_feed(cls, session, user, celery_poller, tasks, url):
        """Subscribe the user to a feed and return the feed ID.

        Raises HTTPError if the feed at the given url can't be subscribed to.
        """
        # add a new feed if it doesn't exist already
        feed = session.query(Feed).filter(Feed.feed_url == url).first()
        if feed is None:
            res = yield celery_poller.run_task(tasks.fetch_feed, url)
            res = yaml.safe_load(res)
            if "error" in res:
                logger.warning("Failed to fetch new feed: '{}'"
                               .format(res['error']))
                raise HTTPError(400, reason=res['error'])
            feed = res['feed']
            existing_feed = session.query(Feed)\
                    .filter(Feed.feed_url == feed.feed_url).all()
            if len(existing_feed) > 0:
                feed = existing_feed[0]
            else:
                # only add the first entry with a given guid
                # this assumes the first one is the newest
                entries_by_guid = {}
                for entry in res['entries']:
                    if entry.guid not in entries_by_guid:
                        entries_by_guid[entry.guid] = entry
                feed.add_all(entries_by_guid.values())
                session.commit()
        if user.has_subscription(feed):
            raise HTTPError(400, reason="Already subscribed to feed")

        # subscribe the user to the feed
        user.subscribe(feed)
        session.commit()
        raise gen.Return(feed.id)

    @asynchronous
    @gen.coroutine
    def post(self):
        """Subscribe the user to a new feed."""
        body = self.require_body_schema({
            'type': 'object',
            'properties': {
                'url': {'type': 'string'},
            },
            'required': ['url'],
        })
        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            feed_id = yield self.subscribe_feed(session, user,
                                                self.celery_poller, self.tasks,
                                                body['url'])
            # TODO: technically this is supposed to be an absolute URL
            self.set_header("Location", "/feeds/{}".format(feed_id))
        self.set_status(201)

class FeedHandler(APIRequestHandler):

    def get(self, feed_id):
        """Return metadata for a subscribed feed."""
        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            feed = session.query(Feed).get(int(feed_id))

            # Make sure the feed exists
            if feed is None:
                raise HTTPError(404, reason='This feed does not exist')

            # Make sure the user is subscribed to this feed
            if not user.has_subscription(feed):
                raise HTTPError(404, reason='This feed does not exist')

            self.write({
                'id': feed.id,
                'name': feed.title,
                'url': feed.site_url,
                'image_url': feed.image_url,
                'unreads': user.get_num_unread_entries(feed),
            })
        self.set_status(200)

    def delete(self, feed_id):
        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            feed = session.query(Feed).get(int(feed_id))

            # Make sure the feed exists
            if feed is None:
                raise HTTPError(404, reason='This feed does not exist')

            # Make sure the user is subscribed to this feed
            if not user.has_subscription(feed):
                raise HTTPError(404, reason='This feed does not exist')

            user.unsubscribe(feed)
        self.set_status(204)

class FeedEntriesHandler(APIRequestHandler):

    def get(self, feed_id):
        """Return a list of entry IDs for a feed."""
        entry_filter = self.get_argument('filter', None)

        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            feed = session.query(Feed).get(int(feed_id))

            # Make sure the feed exists
            if feed is None:
                raise HTTPError(404, reason='This feed does not exist')

            # Make sure the user is subscribed to this feed
            if not user.has_subscription(feed):
                raise HTTPError(404, reason='This feed does not exist')

            # Make sure the filter keyword is valid
            if entry_filter not in ['read', 'unread', None]:
                raise HTTPError(400, reason='Filter keyboard is not valid')

            # Get feed entries
            entries = feed.get_entries(user, entry_filter)
            self.write({'entries': [entry.id for entry in entries]})
            self.set_status(200)

class EntriesHandler(APIRequestHandler):

    @classmethod
    @gen.coroutine
    def truncate(cls, content, length):
        return BeautifulSoup(content).get_text()[:length]

    @asynchronous
    @gen.coroutine
    def get(self, entry_ids):
        entry_ids = [int(entry_id) for entry_id in entry_ids.split(",")]

        try:
            length = int(self.get_argument('truncate', 0))
        except ValueError:
            raise HTTPError(400, reason="Invalid truncation size.")

        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            entries = user.get_entries(entry_ids)

            if len(entries) != len(entry_ids):
                raise HTTPError(404, "Entry does not exist.")

            entries_json = []
            for entry in entries:
                entry_json = {
                    "id": entry.id,
                    "title": entry.title,
                    "pub-date": entry.date,
                    "read": user.has_read(entry),
                    "author": entry.author,
                    "feed_id": entry.feed_id,
                    "url": entry.url,
                }

                if length:
                    entry_json['content'] = yield self.truncate(entry.content,
                                                                length)
                else:
                    entry_json['content'] = entry.content
                entries_json.append(entry_json)

            self.write({'entries': entries_json})
        self.set_status(200)

    def patch(self, entry_ids):
        entry_ids = [int(entry_id) for entry_id in entry_ids.split(",")]
        body = self.require_body_schema({
            "type": "object",
            "properties": {
                "read": {"type": "boolean"},
            },
            "required": ["read"],
        })

        with self.get_db_session() as session:
            user = session.query(User).get(self.require_auth(session))
            entries = user.get_entries(entry_ids)

            if len(entries) != len(entry_ids):
                raise HTTPError(404, "Entry does not exist.")

            if body["read"]:
                for entry in entries:
                    if not user.has_read(entry):
                        user.read(entry)
            else:
                for entry in entries:
                    if user.has_read(entry):
                        user.unread(entry)
        self.set_status(200)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Parses simple newline separated list of IPs.

Docs:
 - https://feodotracker.abuse.ch/blocklist/
 - https://zeustracker.abuse.ch/blocklist.php
"""

import dateutil

from intelmq.lib.bot import ParserBot

SOURCE_FEEDS = {'https://feodotracker.abuse.ch/blocklist/?download=ipblocklist': 'Cridex',
                'https://zeustracker.abuse.ch/blocklist.php?download=ipblocklist': 'Zeus',
                'https://zeustracker.abuse.ch/blocklist.php?download=badips': 'Zeus'}

class AbusechIPParserBot(ParserBot):
    lastgenerated = None

    def parse_line(self, line, report):
        if line.startswith('#'):
            self.tempdata.append(line)
            if 'Generated on' in line:
                row = line.strip('# ')[13:]
                self.lastgenerated = dateutil.parser.parse(row).isoformat()
        else:
            event = self.new_event(report)
            event.add('time.source', self.lastgenerated)
            event.add('classification.taxonomy', 'malicious code')
            event.add('classification.type', 'c&c')
            event.add('source.ip', line)
            event.add("raw", line)
            event.add("malware.name", SOURCE_FEEDS[report["feed.url"]])
            yield event

    def recover_line(self, line):
        return '\n'.join(self.tempdata + [line])

BOT = AbusechIPParserBot
<commit_msg>Parse feeds line independent of fields<commit_after># -*- coding: utf-8 -*-
"""
Parses simple newline separated list of IPs.

Docs:
 - https://feodotracker.abuse.ch/blocklist/
 - https://zeustracker.abuse.ch/blocklist.php
"""
import re

import dateutil

from intelmq.lib.bot import ParserBot
from intelmq.lib import utils
from intelmq.lib.exceptions import PipelineError

FEEDS = {
    'https://feodotracker.abuse.ch/downloads/ipblocklist.csv': {
        'format': [
            'time.source',
            'source.ip',
            'source.port',
            'malware.name'
        ],
        'malware': 'Cridex'
    },
    'https://zeustracker.abuse.ch/blocklist.php?download=ipblocklist': {
        'format': [
            'source.ip'
        ],
        'malware': 'Zeus'
    },
    'https://zeustracker.abuse.ch/blocklist.php?download=badips': {
        'format': [
            'source.ip'
        ],
        'malware': 'Zeus'
    }
}

SOURCE_FEEDS = {'https://feodotracker.abuse.ch/blocklist/?download=ipblocklist': 'Cridex',
                'https://zeustracker.abuse.ch/blocklist.php?download=ipblocklist': 'Zeus',
                'https://zeustracker.abuse.ch/blocklist.php?download=badips': 'Zeus'}

class AbusechIPParserBot(ParserBot):
    lastgenerated = None

    def parse(self, report: dict):
        lines = utils.base64_decode(report.get("raw")).splitlines()
        regex = re.compile(r'^#+.*')  # filter out comments
        top_line = list(filter(lambda i: not regex.search(i), lines))[0]
        feed = report['feed.url']

        fields = top_line.split(',')
        if len(fields) is not len(FEEDS[feed]['format']):
            self.logger.warning(
                f"Feed '{feed}' has not the expected length of fields: {len(fields)} != {len(FEEDS[feed]['format'])}")
            raise PipelineError("Abusech ip parser bad config")

        for line in utils.base64_decode(report.get("raw")).splitlines():
            line = line.strip()
            if not any([line.startswith(prefix) for prefix in self.ignore_lines_starting]):
                yield line

    def parse_line(self, line, report):
        if line.startswith('#'):
            self.tempdata.append(line)
            if 'Generated on' in line:
                row = line.strip('# ')[13:]
                self.lastgenerated = dateutil.parser.parse(row).isoformat()
        else:
            event = self.new_event(report)
            self.__process_feed_fields(event, line, report['feed.url'])
            yield event

    def recover_line(self, line):
        return '\n'.join(self.tempdata + [line])

    def __process_feed_fields(self, event, line, feed):
        event.add('classification.taxonomy', 'malicious code')
        event.add('classification.type', 'c&c')
        event.add('raw', line)

        current_feed_format = FEEDS[feed]['format']
        data = line.split(',') if len(current_feed_format) > 1 else line

        if isinstance(data, list):
            # has multiple fields
            i = 0
            for field in current_feed_format:
                event.add(field, data[i])
                i += 1
        else:
            # has only one 'ip' field
            event.add('time.source', self.lastgenerated)
            event.add('source.ip', line)
            event.add('malware.name', FEEDS[feed]['malware'])

BOT = AbusechIPParserBot
<|endoftext|>


<commit_before>"""

Run demo with
python3 -m interpolation.interpolate interpolate.py
"""
import datetime
import random
import logging
import itertools
import os

import numpy
import pandas

from filter_weather_data.filters import StationRepository
from filter_weather_data import get_repository_parameters
from filter_weather_data import RepositoryParameter

from .interpolator.delaunay_triangulator import DelaunayTriangulator
from .interpolator.nearest_k_finder import NearestKFinder
from . import load_airport
from .interpolator.statistical_interpolator_experimental import get_interpolation_results

class Scorer:
    def __init__(self, target_station_dict, neighbour_station_dicts, start_date, end_date):
        self.target_station_dict = target_station_dict
        self.nearest_k_finder = NearestKFinder(neighbour_station_dicts, start_date, end_date)
        self.delaunay_triangulator = DelaunayTriangulator(neighbour_station_dicts, start_date, end_date)
        self.airport_df = load_airport("EDDH", start_date, end_date)

    def score_nearest_neighbour(self, date, t_actual):
        neighbours = self.nearest_k_finder.find_k_nearest_neighbours(self.target_station_dict, date, 1)
        if len(neighbours) == 1:
            t_nb = neighbours[0][0]
            return (t_nb - t_actual) ** 2
        else:
            return numpy.nan

    def score_airport(self, date, t_actual):
        t_eddh = self.airport_df.loc[date].temperature
        return (t_eddh - t_actual) ** 2

    def score_3_nearest_neighbours(self, date, t_actual):
        relevant_neighbours = self.nearest_k_finder.find_k_nearest_neighbours(self.target_station_dict, date, 3)
        return get_interpolation_results(relevant_neighbours, t_actual, "_cn3")

    def score_5_nearest_neighbours(self, date, t_actual):
        relevant_neighbours = self.nearest_k_finder.find_k_nearest_neighbours(self.target_station_dict, date, 5)
        return get_interpolation_results(relevant_neighbours, t_actual, "_cn5")

    def score_delaunay_neighbours(self, date, t_actual):
        relevant_neighbours = self.delaunay_triangulator.find_delaunay_neighbours(self.target_station_dict, date)
        return get_interpolation_results(relevant_neighbours, t_actual, "_dt")

    def score_all_neighbours(self, date, t_actual):
        relevant_neighbours = self.nearest_k_finder.find_k_nearest_neighbours(self.target_station_dict, date, -1)
        return get_interpolation_results(relevant_neighbours, t_actual, "_all")

def score_interpolation_algorithm_at_date(scorer, date):
    t_actual = scorer.target_station_dict["data_frame"].loc[date].temperature
    results = {
        "nn": scorer.score_nearest_neighbour(date, t_actual),
        "eddh": scorer.score_airport(date, t_actual),
    }
    results.update(scorer.score_3_nearest_neighbours(date, t_actual))
    results.update(scorer.score_5_nearest_neighbours(date, t_actual))
    results.update(scorer.score_all_neighbours(date, t_actual))
    results.update(scorer.score_delaunay_neighbours(date, t_actual))
    return results

def setup_logging(interpolation_name):
    log = logging.getLogger('')

    log.setLevel(logging.DEBUG)
    formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')

    #console_handler = logging.StreamHandler(sys.stdout)
    #console_handler.setFormatter(formatter)
    #log.addHandler(console_handler)

    file_name = "interpolation_{date}_{interpolation_name}.log".format(
        interpolation_name=interpolation_name,
        date=datetime.datetime.now().isoformat().replace(":", "-").replace(".", "-")
    )
    path_to_file_to_log_to = os.path.join(
        os.path.dirname(os.path.realpath(__file__)),
        "log",
        file_name
    )
    file_handler = logging.FileHandler(path_to_file_to_log_to)
    file_handler.setFormatter(formatter)
    log.addHandler(file_handler)

    log.propagate = False

    log.info("### Start new logging")
    return log

def do_interpolation_scoring(
        target_station_dict,
        j,
        target_station_dicts_len,
        neighbour_station_dicts,
        start_date,
        end_date
):
    target_station_name = target_station_dict["name"]
    logging.info("interpolate for " + target_station_name)
    logging.info("currently at " + str(j + 1) + " out of " + target_station_dicts_len)
    logging.info("use " + " ".join([station_dict["name"] for station_dict in neighbour_station_dicts]))

    scorer = Scorer(target_station_dict, neighbour_station_dicts, start_date, end_date)
    scorer.nearest_k_finder.sample_up(target_station_dict, start_date, end_date)
    sum_square_errors = {}
    total_len = len(target_station_dict["data_frame"].index.values)
    each_minute = target_station_dict["data_frame"].index.values
    grouped_by_hour = numpy.array_split(each_minute, total_len / 60)
    each_hour = [numpy.random.choice(hour_group) for hour_group in grouped_by_hour]
    for current_i, date in enumerate(each_hour):
        result = score_interpolation_algorithm_at_date(scorer, date)
        for method, square_error in result.items():
            if method not in sum_square_errors:
                sum_square_errors[method] = {}
                sum_square_errors[method]["total"] = 0
                sum_square_errors[method]["n"] = 0
            if not numpy.isnan(square_error):
                sum_square_errors[method]["total"] += square_error
                sum_square_errors[method]["n"] += 1

    for method, result in sum_square_errors.items():
        if sum_square_errors[method]["n"] > 0:
            method_rmse = numpy.sqrt(sum_square_errors[method]["total"] / sum_square_errors[method]["n"])
        else:
            method_rmse = numpy.nan
        sum_square_errors[method]["rmse"] = method_rmse
        score_str = "%.3f" % method_rmse
        logging.info(method + " " * (12 - len(method)) + score_str + " n=" + str(sum_square_errors[method]["n"]))

    logging.info("end method list")

    data_dict = {}
    for method in sum_square_errors.keys():
        data_dict[method + "--rmse"] = [sum_square_errors[method]["rmse"]]
        data_dict[method + "--n"] = [sum_square_errors[method]["n"]]
        data_dict[method + "--total"] = [sum_square_errors[method]["total"]]
    return pandas.DataFrame(data=data_dict)

def score_algorithm(start_date, end_date, repository_parameters, limit=0, interpolation_name="NONE"):
    station_repository = StationRepository(*repository_parameters)
    station_dicts = station_repository.load_all_stations(start_date, end_date, limit=limit)

    # separate in two sets
    random.shuffle(station_dicts)
    separator = int(.3 * len(station_dicts))  # 70% vs 30%
    target_station_dicts, neighbour_station_dicts = station_dicts[:separator], station_dicts[separator:]

    setup_logging(interpolation_name)
    logging.info("General Overview")
    logging.info("targets: " + " ".join([station_dict["name"] for station_dict in target_station_dicts]))
    logging.info("neighbours: " + " ".join([station_dict["name"] for station_dict in neighbour_station_dicts]))
    logging.info("End overview")

    logging.info("Several Runs")
    target_station_dicts_len = str(len(target_station_dicts))

    overall_result = itertools.starmap(do_interpolation_scoring, [
        [
            target_station_dict,
            j,
            target_station_dicts_len,
            neighbour_station_dicts,
            start_date,
            end_date
        ] for j, target_station_dict in enumerate(target_station_dicts)
    ])

    logging.info("end targets")

    logging.info("overall results")
    overall_result_df = pandas.concat(overall_result)
    column_names = overall_result_df.columns.values.tolist()
    methods = set()
    for column_name in column_names:
        method, value = column_name.split("--")
        methods.update([method])
    for method in methods:
        overall_total = numpy.nansum(overall_result_df[method + "--total"])
        overall_n = int(numpy.nansum(overall_result_df[method + "--n"]))
        overall_rmse = numpy.sqrt(overall_total / overall_n)
        score_str = "%.5f" % overall_rmse
        logging.info(method + " " * (12 - len(method)) + score_str + " n=" + str(overall_n))

    overall_result_df.to_csv("interpolation_result_{date}_{interpolation_name}.csv".format(
        date=datetime.datetime.now().isoformat().replace(":", "-").replace(".", "-"),
        interpolation_name=interpolation_name
    ))

def demo():
    start_date = "2016-01-31"
    end_date = "2016-12-31"  # "2016-02-01"
    repository_parameters = get_repository_parameters(RepositoryParameter.ONLY_OUTDOOR_AND_SHADED)
    score_algorithm(start_date, end_date, repository_parameters, limit=60, interpolation_name="test")

if __name__ == "__main__":
    demo()
<commit_msg>speed up by removing unnecessary tasks.<commit_after>"""

Run demo with
python3 -m interpolation.interpolate interpolate.py
"""
import datetime
import random
import logging
import itertools
import os

import numpy
import pandas

from filter_weather_data.filters import StationRepository
from filter_weather_data import get_repository_parameters
from filter_weather_data import RepositoryParameter

from .interpolator.nearest_k_finder import NearestKFinder
from .interpolator.statistical_interpolator_experimental import get_interpolation_results

class Scorer:
    def __init__(self, target_station_dict, neighbour_station_dicts, start_date, end_date):
        self.target_station_dict = target_station_dict
        self.nearest_k_finder = NearestKFinder(neighbour_station_dicts, start_date, end_date)

    def score_3_nearest_neighbours(self, date, t_actual):
        relevant_neighbours = self.nearest_k_finder.find_k_nearest_neighbours(self.target_station_dict, date, 3)
        return get_interpolation_results(relevant_neighbours, t_actual, "_cn3")

    def score_5_nearest_neighbours(self, date, t_actual):
        relevant_neighbours = self.nearest_k_finder.find_k_nearest_neighbours(self.target_station_dict, date, 5)
        return get_interpolation_results(relevant_neighbours, t_actual, "_cn5")

    def score_all_neighbours(self, date, t_actual):
        relevant_neighbours = self.nearest_k_finder.find_k_nearest_neighbours(self.target_station_dict, date, -1)
        return get_interpolation_results(relevant_neighbours, t_actual, "_all")

def score_interpolation_algorithm_at_date(scorer, date):
    t_actual = scorer.target_station_dict["data_frame"].loc[date].temperature
    results = {}
    results.update(scorer.score_3_nearest_neighbours(date, t_actual))
    results.update(scorer.score_5_nearest_neighbours(date, t_actual))
    results.update(scorer.score_all_neighbours(date, t_actual))
    return results

def setup_logging(interpolation_name):
    log = logging.getLogger('')

    log.setLevel(logging.DEBUG)
    formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')

    #console_handler = logging.StreamHandler(sys.stdout)
    #console_handler.setFormatter(formatter)
    #log.addHandler(console_handler)

    file_name = "interpolation_{date}_{interpolation_name}.log".format(
        interpolation_name=interpolation_name,
        date=datetime.datetime.now().isoformat().replace(":", "-").replace(".", "-")
    )
    path_to_file_to_log_to = os.path.join(
        os.path.dirname(os.path.realpath(__file__)),
        "log",
        file_name
    )
    file_handler = logging.FileHandler(path_to_file_to_log_to)
    file_handler.setFormatter(formatter)
    log.addHandler(file_handler)

    log.propagate = False

    log.info("### Start new logging")
    return log

def do_interpolation_scoring(
        target_station_dict,
        j,
        target_station_dicts_len,
        neighbour_station_dicts,
        start_date,
        end_date
):
    target_station_name = target_station_dict["name"]
    logging.info("interpolate for " + target_station_name)
    logging.info("currently at " + str(j + 1) + " out of " + target_station_dicts_len)
    logging.info("use " + " ".join([station_dict["name"] for station_dict in neighbour_station_dicts]))

    scorer = Scorer(target_station_dict, neighbour_station_dicts, start_date, end_date)
    scorer.nearest_k_finder.sample_up(target_station_dict, start_date, end_date)
    sum_square_errors = {}
    total_len = len(target_station_dict["data_frame"].index.values)
    each_minute = target_station_dict["data_frame"].index.values
    grouped_by_hour = numpy.array_split(each_minute, total_len / 60)
    each_hour = [numpy.random.choice(hour_group) for hour_group in grouped_by_hour]
    for current_i, date in enumerate(each_hour):
        result = score_interpolation_algorithm_at_date(scorer, date)
        for method, square_error in result.items():
            if method not in sum_square_errors:
                sum_square_errors[method] = {}
                sum_square_errors[method]["total"] = 0
                sum_square_errors[method]["n"] = 0
            if not numpy.isnan(square_error):
                sum_square_errors[method]["total"] += square_error
                sum_square_errors[method]["n"] += 1

    for method, result in sum_square_errors.items():
        if sum_square_errors[method]["n"] > 0:
            method_rmse = numpy.sqrt(sum_square_errors[method]["total"] / sum_square_errors[method]["n"])
        else:
            method_rmse = numpy.nan
        sum_square_errors[method]["rmse"] = method_rmse
        score_str = "%.3f" % method_rmse
        logging.info(method + " " * (12 - len(method)) + score_str + " n=" + str(sum_square_errors[method]["n"]))

    logging.info("end method list")

    data_dict = {}
    for method in sum_square_errors.keys():
        data_dict[method + "--rmse"] = [sum_square_errors[method]["rmse"]]
        data_dict[method + "--n"] = [sum_square_errors[method]["n"]]
        data_dict[method + "--total"] = [sum_square_errors[method]["total"]]
    return pandas.DataFrame(data=data_dict)

def score_algorithm(start_date, end_date, repository_parameters, limit=0, interpolation_name="NONE"):
    station_repository = StationRepository(*repository_parameters)
    station_dicts = station_repository.load_all_stations(start_date, end_date, limit=limit)

    # separate in two sets
    random.shuffle(station_dicts)
    separator = int(.3 * len(station_dicts))  # 70% vs 30%
    target_station_dicts, neighbour_station_dicts = station_dicts[:separator], station_dicts[separator:]

    setup_logging(interpolation_name)
    logging.info("General Overview")
    logging.info("targets: " + " ".join([station_dict["name"] for station_dict in target_station_dicts]))
    logging.info("neighbours: " + " ".join([station_dict["name"] for station_dict in neighbour_station_dicts]))
    logging.info("End overview")

    logging.info("Several Runs")
    target_station_dicts_len = str(len(target_station_dicts))

    overall_result = itertools.starmap(do_interpolation_scoring, [
        [
            target_station_dict,
            j,
            target_station_dicts_len,
            neighbour_station_dicts,
            start_date,
            end_date
        ] for j, target_station_dict in enumerate(target_station_dicts)
    ])

    logging.info("end targets")

    logging.info("overall results")
    overall_result_df = pandas.concat(overall_result)
    column_names = overall_result_df.columns.values.tolist()
    methods = set()
    for column_name in column_names:
        method, value = column_name.split("--")
        methods.update([method])
    for method in methods:
        overall_total = numpy.nansum(overall_result_df[method + "--total"])
        overall_n = int(numpy.nansum(overall_result_df[method + "--n"]))
        overall_rmse = numpy.sqrt(overall_total / overall_n)
        score_str = "%.5f" % overall_rmse
        logging.info(method + " " * (12 - len(method)) + score_str + " n=" + str(overall_n))

    overall_result_df.to_csv("interpolation_result_{date}_{interpolation_name}.csv".format(
        date=datetime.datetime.now().isoformat().replace(":", "-").replace(".", "-"),
        interpolation_name=interpolation_name
    ))

def demo():
    start_date = "2016-01-31"
    end_date = "2016-12-31"  # "2016-02-01"
    repository_parameters = get_repository_parameters(RepositoryParameter.ONLY_OUTDOOR_AND_SHADED)
    score_algorithm(start_date, end_date, repository_parameters, limit=60, interpolation_name="test")

if __name__ == "__main__":
    demo()
<|endoftext|>