<commit_before># coding: utf-8
from __future__ import absolute_import, unicode_literals

import logging

from celery import shared_task

from nodeconductor.core import models as core_models
from nodeconductor.core.models import SynchronizationStates
from nodeconductor.core.tasks import tracked_processing, set_state, StateChangeError
from nodeconductor.core.log import EventLoggerAdapter
from nodeconductor.iaas import models
from nodeconductor.iaas.backend import CloudBackendError
from nodeconductor.monitoring.zabbix.api_client import ZabbixApiClient
from nodeconductor.monitoring.zabbix.errors import ZabbixError

logger = logging.getLogger(__name__)
event_logger = EventLoggerAdapter(logger)

# XXX: There are no usages of this error.
class ResizingError(KeyError, models.Instance.DoesNotExist):
    pass

def create_zabbix_host_and_service(instance, warn_if_exists=True):
    try:
        zabbix_client = ZabbixApiClient()
        zabbix_client.create_host(instance, warn_if_host_exists=warn_if_exists)
        zabbix_client.create_service(instance, warn_if_service_exists=warn_if_exists)
    except Exception as e:
        # task does not have to fail if something is wrong with zabbix
        logger.error('Zabbix host creation flow has broken %s' % e, exc_info=1)
        event_logger.error(
            'Zabbix host creation flow has broken %s', e,
            extra={'instance': instance, 'event_type': 'zabbix_host_creation'}
        )

def delete_zabbix_host_and_service(instance):
    try:
        zabbix_client = ZabbixApiClient()
        zabbix_client.delete_host(instance)
        zabbix_client.delete_service(instance)
    except ZabbixError as e:
        # task does not have to fail if something is wrong with zabbix
        logger.error('Zabbix host deletion flow has broken %s' % e, exc_info=1)
        event_logger.error(
            'Zabbix host deletion flow has broken %s', e,
            extra={'instance': instance, 'event_type': 'zabbix_host_deletion'}
        )

@shared_task
@tracked_processing(models.Instance, processing_state='begin_provisioning', desired_state='set_online')
def schedule_provisioning(instance_uuid, backend_flavor_id, system_volume_id=None, data_volume_id=None):
    instance = models.Instance.objects.get(uuid=instance_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    try:
        backend.provision_instance(instance, backend_flavor_id, system_volume_id, data_volume_id)
    finally:
        # the function below should never fail
        create_zabbix_host_and_service(instance)

@shared_task
@tracked_processing(models.Instance, processing_state='begin_stopping', desired_state='set_offline')
def schedule_stopping(instance_uuid):
    instance = models.Instance.objects.get(uuid=instance_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    backend.stop_instance(instance)

@shared_task
@tracked_processing(models.Instance, processing_state='begin_restarting', desired_state='set_restarted')
def schedule_restarting(instance_uuid):
    instance = models.Instance.objects.get(uuid=instance_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    backend.restart_instance(instance)

@shared_task
@tracked_processing(models.Instance, processing_state='begin_starting', desired_state='set_online')
def schedule_starting(instance_uuid):
    instance = models.Instance.objects.get(uuid=instance_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    backend.start_instance(instance)

@shared_task
def schedule_deleting(instance_uuid):
    try:
        set_state(models.Instance, instance_uuid, 'begin_deleting')
    except StateChangeError:
        # No logging is needed since set_state already logged everything
        return

    # noinspection PyBroadException
    try:
        instance = models.Instance.objects.get(uuid=instance_uuid)

        backend = instance.cloud_project_membership.cloud.get_backend()
        backend.delete_instance(instance)

        delete_zabbix_host_and_service(instance)
    except Exception:
        # noinspection PyProtectedMember
        logger.exception('Failed to begin_deleting Instance with id %s', instance_uuid)
        event_logger.exception(
            'Failed to begin_deleting Instance with id %s', instance_uuid,
            extra={'instance': instance, 'event_type': 'instance_deletion'}
        )
        set_state(models.Instance, instance_uuid, 'set_erred')
    else:
        # Actually remove the instance from the database
        models.Instance.objects.filter(uuid=instance_uuid).delete()

@shared_task
@tracked_processing(models.Instance, processing_state='begin_resizing', desired_state='set_offline')
def update_flavor(instance_uuid, flavor_uuid):
    instance = models.Instance.objects.get(uuid=instance_uuid)
    flavor = models.Flavor.objects.get(uuid=flavor_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    backend.update_flavor(instance, flavor)

@shared_task
@tracked_processing(models.Instance, processing_state='begin_resizing', desired_state='set_offline')
def extend_disk(instance_uuid):
    instance = models.Instance.objects.get(uuid=instance_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    backend.extend_disk(instance)

@shared_task
def push_instance_security_groups(instance_uuid):
    instance = models.Instance.objects.get(uuid=instance_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    backend.push_instance_security_groups(instance)

@shared_task
@tracked_processing(
    models.Cloud,
    processing_state='begin_syncing',
    desired_state='set_in_sync',
)
def pull_cloud_account(cloud_account_uuid):
    cloud_account = models.Cloud.objects.get(uuid=cloud_account_uuid)

    backend = cloud_account.get_backend()
    backend.pull_cloud_account(cloud_account)

@shared_task
def pull_cloud_accounts():
    # TODO: Extract to a service
    queryset = models.Cloud.objects.filter(state=SynchronizationStates.IN_SYNC)

    for cloud_account in queryset.iterator():
        cloud_account.schedule_syncing()
        cloud_account.save()

        pull_cloud_account.delay(cloud_account.uuid.hex)

@shared_task
@tracked_processing(
    models.Cloud,
    processing_state='begin_syncing',
    desired_state='set_in_sync',
)
def sync_cloud_account(cloud_account_uuid):
    cloud = models.Cloud.objects.get(uuid=cloud_account_uuid)

    backend = cloud.get_backend()
    backend.push_cloud_account(cloud)
    backend.pull_cloud_account(cloud)

@shared_task
@tracked_processing(
    models.CloudProjectMembership,
    processing_state='begin_syncing',
    desired_state='set_in_sync',
)
def pull_cloud_membership(membership_pk):
    membership = models.CloudProjectMembership.objects.get(pk=membership_pk)

    backend = membership.cloud.get_backend()
    backend.pull_security_groups(membership)
    backend.pull_instances(membership)
    backend.pull_resource_quota(membership)
    backend.pull_resource_quota_usage(membership)
    backend.pull_floating_ips(membership)

@shared_task
def pull_cloud_memberships():
    # TODO: Extract to a service
    queryset = models.CloudProjectMembership.objects.filter(state=SynchronizationStates.IN_SYNC)

    for membership in queryset.iterator():
        membership.schedule_syncing()
        membership.save()

        pull_cloud_membership.delay(membership.pk)

@shared_task
@tracked_processing(
    models.CloudProjectMembership,
    processing_state='begin_syncing',
    desired_state='set_in_sync',
)
def sync_cloud_membership(membership_pk):
    membership = models.CloudProjectMembership.objects.get(pk=membership_pk)

    backend = membership.cloud.get_backend()

    # Propagate cloud-project membership itself
    backend.push_membership(membership)

    # Propagate ssh public keys of users involved in the project
    for public_key in core_models.SshPublicKey.objects.filter(
            user__groups__projectrole__project=membership.project).iterator():
        try:
            backend.push_ssh_public_key(membership, public_key)
        except CloudBackendError:
            logger.warn(
                'Failed to push public key %s to cloud membership %s',
                public_key.uuid, membership.pk,
                exc_info=1,
            )
            event_logger.warning(
                'Failed to push public key %s to cloud membership %s',
                public_key.uuid, membership.pk,
                extra={'project': membership.project, 'cloud': membership.cloud, 'event_type': 'sync_cloud_membership'}
            )

    # Propagate membership security groups
    try:
        backend.push_security_groups(membership)
    except CloudBackendError:
        logger.warn(
            'Failed to push security groups to cloud membership %s',
            membership.pk,
            exc_info=1,
        )
        event_logger.warning(
            'Failed to push security groups to cloud membership %s',
            public_key.uuid, membership.pk,
            extra={'project': membership.project, 'cloud': membership.cloud, 'event_type': 'sync_cloud_membership'}
        )

    # Pull created membership quotas
    try:
        backend.pull_resource_quota(membership)
        backend.pull_resource_quota_usage(membership)
    except CloudBackendError:
        logger.warn(
            'Failed to pull resource quota and usage data to cloud membership %s',
            membership.pk,
            exc_info=1,
        )

@shared_task
@tracked_processing(
    models.Cloud,
    processing_state='begin_syncing',
    desired_state='set_in_sync',
)
def pull_images(cloud_account_uuid):
    cloud = models.Cloud.objects.get(uuid=cloud_account_uuid)
    cloud.get_backend().pull_images(cloud)

@shared_task
def push_ssh_public_keys(ssh_public_keys_uuids, membership_pks):
    public_keys = core_models.SshPublicKey.objects.filter(uuid__in=ssh_public_keys_uuids)

    existing_keys = set(k.uuid.hex for k in public_keys)
    missing_keys = set(ssh_public_keys_uuids) - existing_keys
    if missing_keys:
        logging.warn(
            'Failed to push missing public keys: %s',
            ', '.join(missing_keys)
        )

    membership_queryset = models.CloudProjectMembership.objects.filter(
        pk__in=membership_pks)

    potential_rerunnable = []
    for membership in membership_queryset.iterator():
        if membership.state != core_models.SynchronizationStates.IN_SYNC:
            logging.warn(
                'Not pushing public keys to cloud membership %s which is in state %s. Re-scheduling.',
                membership.pk, membership.get_state_display()
            )
            if membership.state != core_models.SynchronizationStates.ERRED:
                # reschedule a task for this membership if membership is in a sane state
                potential_rerunnable.append(membership.id)
            continue

        backend = membership.cloud.get_backend()
        for public_key in public_keys:
            try:
                backend.push_ssh_public_key(membership, public_key)
            except CloudBackendError:
                logger.warn(
                    'Failed to push public key %s to cloud membership %s',
                    public_key.uuid, membership.pk,
                    exc_info=1,
                )
    # reschedule sync to membership that were blocked
    # TODO: temporarily disable -- NC-305
    #push_ssh_public_keys.delay(ssh_public_keys_uuids, potential_rerunnable)

@shared_task
def check_cloud_memberships_quotas():
    threshold = 0.80  # Could have been configurable...

    resources = models.AbstractResourceQuota._meta.get_all_field_names()

    queryset = (
        models.CloudProjectMembership.objects
        .all()
        .select_related(
            'cloud',
            'cloud__customer',
            'project',
            'resource_quota',
            'resource_quota_usage',
        )
        .prefetch_related(
            'project__project_groups',
        )
    )

    for membership in queryset.iterator():
        quota = membership.resource_quota
        usage = membership.resource_quota_usage

        for resource_name in resources:
            resource_quota = getattr(quota, resource_name)
            resource_usage = getattr(usage, resource_name)

            if resource_usage >= threshold * resource_quota:
                event_logger.warning(
                    "%s quota limit is about to be reached", resource_name,
                    extra=dict(
                        event_type='quota_threshold_reached',
                        cloud=membership.cloud,
                        project=membership.project,
                        project_group=membership.project.project_groups.first(),
                    ))

@shared_task
def sync_instances_with_zabbix():
    for instance in models.Instance.objects.all():
        if instance.backend_id:
            logger.info(
                'Synchronizing instance %s with zabbix',
                instance.uuid,
                exc_info=1,
            )
            create_zabbix_host_and_service(instance, warn_if_exists=False)
<commit_msg>More resilient checks for quota usage<commit_after># coding: utf-8
from __future__ import absolute_import, unicode_literals

import logging

from celery import shared_task
from django.core.exceptions import ObjectDoesNotExist

from nodeconductor.core import models as core_models
from nodeconductor.core.models import SynchronizationStates
from nodeconductor.core.tasks import tracked_processing, set_state, StateChangeError
from nodeconductor.core.log import EventLoggerAdapter
from nodeconductor.iaas import models
from nodeconductor.iaas.backend import CloudBackendError
from nodeconductor.monitoring.zabbix.api_client import ZabbixApiClient
from nodeconductor.monitoring.zabbix.errors import ZabbixError

logger = logging.getLogger(__name__)
event_logger = EventLoggerAdapter(logger)

# XXX: There are no usages of this error.
class ResizingError(KeyError, models.Instance.DoesNotExist):
    pass

def create_zabbix_host_and_service(instance, warn_if_exists=True):
    try:
        zabbix_client = ZabbixApiClient()
        zabbix_client.create_host(instance, warn_if_host_exists=warn_if_exists)
        zabbix_client.create_service(instance, warn_if_service_exists=warn_if_exists)
    except Exception as e:
        # task does not have to fail if something is wrong with zabbix
        logger.error('Zabbix host creation flow has broken %s' % e, exc_info=1)
        event_logger.error(
            'Zabbix host creation flow has broken %s', e,
            extra={'instance': instance, 'event_type': 'zabbix_host_creation'}
        )

def delete_zabbix_host_and_service(instance):
    try:
        zabbix_client = ZabbixApiClient()
        zabbix_client.delete_host(instance)
        zabbix_client.delete_service(instance)
    except ZabbixError as e:
        # task does not have to fail if something is wrong with zabbix
        logger.error('Zabbix host deletion flow has broken %s' % e, exc_info=1)
        event_logger.error(
            'Zabbix host deletion flow has broken %s', e,
            extra={'instance': instance, 'event_type': 'zabbix_host_deletion'}
        )

@shared_task
@tracked_processing(models.Instance, processing_state='begin_provisioning', desired_state='set_online')
def schedule_provisioning(instance_uuid, backend_flavor_id, system_volume_id=None, data_volume_id=None):
    instance = models.Instance.objects.get(uuid=instance_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    try:
        backend.provision_instance(instance, backend_flavor_id, system_volume_id, data_volume_id)
    finally:
        # the function below should never fail
        create_zabbix_host_and_service(instance)

@shared_task
@tracked_processing(models.Instance, processing_state='begin_stopping', desired_state='set_offline')
def schedule_stopping(instance_uuid):
    instance = models.Instance.objects.get(uuid=instance_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    backend.stop_instance(instance)

@shared_task
@tracked_processing(models.Instance, processing_state='begin_restarting', desired_state='set_restarted')
def schedule_restarting(instance_uuid):
    instance = models.Instance.objects.get(uuid=instance_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    backend.restart_instance(instance)

@shared_task
@tracked_processing(models.Instance, processing_state='begin_starting', desired_state='set_online')
def schedule_starting(instance_uuid):
    instance = models.Instance.objects.get(uuid=instance_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    backend.start_instance(instance)

@shared_task
def schedule_deleting(instance_uuid):
    try:
        set_state(models.Instance, instance_uuid, 'begin_deleting')
    except StateChangeError:
        # No logging is needed since set_state already logged everything
        return

    # noinspection PyBroadException
    try:
        instance = models.Instance.objects.get(uuid=instance_uuid)

        backend = instance.cloud_project_membership.cloud.get_backend()
        backend.delete_instance(instance)

        delete_zabbix_host_and_service(instance)
    except Exception:
        # noinspection PyProtectedMember
        logger.exception('Failed to begin_deleting Instance with id %s', instance_uuid)
        event_logger.exception(
            'Failed to begin_deleting Instance with id %s', instance_uuid,
            extra={'instance': instance, 'event_type': 'instance_deletion'}
        )
        set_state(models.Instance, instance_uuid, 'set_erred')
    else:
        # Actually remove the instance from the database
        models.Instance.objects.filter(uuid=instance_uuid).delete()

@shared_task
@tracked_processing(models.Instance, processing_state='begin_resizing', desired_state='set_offline')
def update_flavor(instance_uuid, flavor_uuid):
    instance = models.Instance.objects.get(uuid=instance_uuid)
    flavor = models.Flavor.objects.get(uuid=flavor_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    backend.update_flavor(instance, flavor)

@shared_task
@tracked_processing(models.Instance, processing_state='begin_resizing', desired_state='set_offline')
def extend_disk(instance_uuid):
    instance = models.Instance.objects.get(uuid=instance_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    backend.extend_disk(instance)

@shared_task
def push_instance_security_groups(instance_uuid):
    instance = models.Instance.objects.get(uuid=instance_uuid)

    backend = instance.cloud_project_membership.cloud.get_backend()
    backend.push_instance_security_groups(instance)

@shared_task
@tracked_processing(
    models.Cloud,
    processing_state='begin_syncing',
    desired_state='set_in_sync',
)
def pull_cloud_account(cloud_account_uuid):
    cloud_account = models.Cloud.objects.get(uuid=cloud_account_uuid)

    backend = cloud_account.get_backend()
    backend.pull_cloud_account(cloud_account)

@shared_task
def pull_cloud_accounts():
    # TODO: Extract to a service
    queryset = models.Cloud.objects.filter(state=SynchronizationStates.IN_SYNC)

    for cloud_account in queryset.iterator():
        cloud_account.schedule_syncing()
        cloud_account.save()

        pull_cloud_account.delay(cloud_account.uuid.hex)

@shared_task
@tracked_processing(
    models.Cloud,
    processing_state='begin_syncing',
    desired_state='set_in_sync',
)
def sync_cloud_account(cloud_account_uuid):
    cloud = models.Cloud.objects.get(uuid=cloud_account_uuid)

    backend = cloud.get_backend()
    backend.push_cloud_account(cloud)
    backend.pull_cloud_account(cloud)

@shared_task
@tracked_processing(
    models.CloudProjectMembership,
    processing_state='begin_syncing',
    desired_state='set_in_sync',
)
def pull_cloud_membership(membership_pk):
    membership = models.CloudProjectMembership.objects.get(pk=membership_pk)

    backend = membership.cloud.get_backend()
    backend.pull_security_groups(membership)
    backend.pull_instances(membership)
    backend.pull_resource_quota(membership)
    backend.pull_resource_quota_usage(membership)
    backend.pull_floating_ips(membership)

@shared_task
def pull_cloud_memberships():
    # TODO: Extract to a service
    queryset = models.CloudProjectMembership.objects.filter(state=SynchronizationStates.IN_SYNC)

    for membership in queryset.iterator():
        membership.schedule_syncing()
        membership.save()

        pull_cloud_membership.delay(membership.pk)

@shared_task
@tracked_processing(
    models.CloudProjectMembership,
    processing_state='begin_syncing',
    desired_state='set_in_sync',
)
def sync_cloud_membership(membership_pk):
    membership = models.CloudProjectMembership.objects.get(pk=membership_pk)

    backend = membership.cloud.get_backend()

    # Propagate cloud-project membership itself
    backend.push_membership(membership)

    # Propagate ssh public keys of users involved in the project
    for public_key in core_models.SshPublicKey.objects.filter(
            user__groups__projectrole__project=membership.project).iterator():
        try:
            backend.push_ssh_public_key(membership, public_key)
        except CloudBackendError:
            logger.warn(
                'Failed to push public key %s to cloud membership %s',
                public_key.uuid, membership.pk,
                exc_info=1,
            )
            event_logger.warning(
                'Failed to push public key %s to cloud membership %s',
                public_key.uuid, membership.pk,
                extra={'project': membership.project, 'cloud': membership.cloud, 'event_type': 'sync_cloud_membership'}
            )

    # Propagate membership security groups
    try:
        backend.push_security_groups(membership)
    except CloudBackendError:
        logger.warn(
            'Failed to push security groups to cloud membership %s',
            membership.pk,
            exc_info=1,
        )
        event_logger.warning(
            'Failed to push security groups to cloud membership %s',
            public_key.uuid, membership.pk,
            extra={'project': membership.project, 'cloud': membership.cloud, 'event_type': 'sync_cloud_membership'}
        )

    # Pull created membership quotas
    try:
        backend.pull_resource_quota(membership)
        backend.pull_resource_quota_usage(membership)
    except CloudBackendError:
        logger.warn(
            'Failed to pull resource quota and usage data to cloud membership %s',
            membership.pk,
            exc_info=1,
        )

@shared_task
@tracked_processing(
    models.Cloud,
    processing_state='begin_syncing',
    desired_state='set_in_sync',
)
def pull_images(cloud_account_uuid):
    cloud = models.Cloud.objects.get(uuid=cloud_account_uuid)
    cloud.get_backend().pull_images(cloud)

@shared_task
def push_ssh_public_keys(ssh_public_keys_uuids, membership_pks):
    public_keys = core_models.SshPublicKey.objects.filter(uuid__in=ssh_public_keys_uuids)

    existing_keys = set(k.uuid.hex for k in public_keys)
    missing_keys = set(ssh_public_keys_uuids) - existing_keys
    if missing_keys:
        logging.warn(
            'Failed to push missing public keys: %s',
            ', '.join(missing_keys)
        )

    membership_queryset = models.CloudProjectMembership.objects.filter(
        pk__in=membership_pks)

    potential_rerunnable = []
    for membership in membership_queryset.iterator():
        if membership.state != core_models.SynchronizationStates.IN_SYNC:
            logging.warn(
                'Not pushing public keys to cloud membership %s which is in state %s. Re-scheduling.',
                membership.pk, membership.get_state_display()
            )
            if membership.state != core_models.SynchronizationStates.ERRED:
                # reschedule a task for this membership if membership is in a sane state
                potential_rerunnable.append(membership.id)
            continue

        backend = membership.cloud.get_backend()
        for public_key in public_keys:
            try:
                backend.push_ssh_public_key(membership, public_key)
            except CloudBackendError:
                logger.warn(
                    'Failed to push public key %s to cloud membership %s',
                    public_key.uuid, membership.pk,
                    exc_info=1,
                )
    # reschedule sync to membership that were blocked
    # TODO: temporarily disable -- NC-305
    #push_ssh_public_keys.delay(ssh_public_keys_uuids, potential_rerunnable)

@shared_task
def check_cloud_memberships_quotas():
    threshold = 0.80  # Could have been configurable...

    resources = models.AbstractResourceQuota._meta.get_all_field_names()

    queryset = (
        models.CloudProjectMembership.objects
        .all()
        .select_related(
            'cloud',
            'cloud__customer',
            'project',
            'resource_quota',
            'resource_quota_usage',
        )
        .prefetch_related(
            'project__project_groups',
        )
    )

    for membership in queryset.iterator():
        try:
            quota = membership.resource_quota
            usage = membership.resource_quota_usage
        except ObjectDoesNotExist as e:
            logger.exception('Missing quota or usage')
            continue

        for resource_name in resources:
            resource_quota = getattr(quota, resource_name)
            resource_usage = getattr(usage, resource_name)

            if resource_usage >= threshold * resource_quota:
                event_logger.warning(
                    "%s quota limit is about to be reached", resource_name,
                    extra=dict(
                        event_type='quota_threshold_reached',
                        cloud=membership.cloud,
                        project=membership.project,
                        project_group=membership.project.project_groups.first(),
                    ))

@shared_task
def sync_instances_with_zabbix():
    for instance in models.Instance.objects.all():
        if instance.backend_id:
            logger.info(
                'Synchronizing instance %s with zabbix',
                instance.uuid,
                exc_info=1,
            )
            create_zabbix_host_and_service(instance, warn_if_exists=False)
<|endoftext|>


<commit_before>#! /usr/bin/python

# PyGreen
# Copyright (c) 2013, Nicolas Vanhoren
#
# Released under the MIT license
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
# Software, and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN
# AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

from __future__ import unicode_literals, print_function

import bottle
import os.path
from mako.lookup import TemplateLookup
import os, os.path, itertools
import wsgiref.handlers
import sys, logging, re
import argparse
import sys
import markdown
import waitress
import utils

_logger = logging.getLogger(__name__)

# always adds the location of the default templates

op = os.path
dn = op.dirname

class File(object):
    "Represents a file object within PyBlue"

    def __init__(self, fname, root):
        self.root  = root
        self.fname = fname
        self.fpath = os.path.join(root, fname)
        self.dname = dn(self.fpath)
        self.ext   = os.path.splitext(fname)[1]

        self.meta  =  dict(name=self.nice_name, sortkey="2", tags=set("data"))
        if not self.skip_file:
            self.meta.update(utils.parse_meta(self.fpath))

    @property
    def nice_name(self):
        "Attempts to generate a nicer name from the filename"
        head, tail = os.path.split(self.fname)
        base, ext = os.path.splitext(tail)
        name = base.title().replace("-", " ").replace("_", " ")
        if self.is_image:
            name = name + self.ext
        return name

    def content(self):
        return file(self.fpath).read()

    @property
    def is_image(self):
        return self.ext in (".png", ".jpg", ".gif")

    @property
    def name(self):
        name = self.meta.get("name", "*** name not set ***")
        return name

    @property
    def sortkey(self):
        return self.meta.get("sortkey", "2")

    @property
    def size(self):
        "File size in KB"
        return utils.get_size(self.fpath)

    @property
    def skip_file(self):
        return self.size > utils.MAX_SIZE_KB

    def write(self, output_folder, text):
        loc = os.path.join(output_folder, self.fname)
        if op.abspath(loc) == op.abspath(self.fpath):
            raise Exception("may not overwrite the original file %s" % loc)

        d = os.path.dirname(loc)
        if not os.path.exists(d):
            os.makedirs(d)
        with open(loc, "wb") as fp:
            fp.write(text)

    def url(self, start=None):

        start = start or self

        rpath = op.relpath(self.root, start.dname)
        rpath = op.join(rpath, self.fname)

        return rpath, self.name

    def __repr__(self):
        return "File: %s (%s)" % (self.name, self.fname)

class PyGreen:
    TEMPLATE_DIR = op.abspath(op.join(op.split(__file__)[0], "templates"))

    def __init__(self):

        # the Bottle application
        self.app = bottle.Bottle()
        # a set of strings that identifies the extension of the files
        # that should be processed using Mako
        self.template_exts = set(["html"])
        # the folder where the files to serve are located. Do not set
        # directly, use set_folder instead
        self.folder = "."

        # recrawls directories on each request, use in serving
        self.refresh = True

        # the collection of files
        self.files = None

        # the TemplateLookup of Mako
        self.templates = TemplateLookup(directories=[self.folder, self.TEMPLATE_DIR],
                                        imports=["from markdown import markdown"],
                                        input_encoding='iso-8859-1',
                                        collection_size=100,
        )
        # A list of regular expression. Files whose the name match
        # one of those regular expressions will not be outputed when generating
        # a static version of the web site
        self.file_exclusion = [r".*\.mako", r".*\.py", r"(^|.*\/)\..*"]

        def is_public(path):
            for ex in self.file_exclusion:
                if re.match(ex, path):
                    return False
            return True

        def base_lister():
            files = []
            for dirpath, dirnames, filenames in os.walk(self.folder):
                filenames.sort()
                for f in filenames:
                    absp = os.path.join(dirpath, f)
                    path = os.path.relpath(absp, self.folder)
                    if is_public(path):
                        files.append(path)
            return files

            # A list of function. Each function must return a list of paths

        # of files to export during the generation of the static web site.
        # The default one simply returns all the files contained in the folder.
        # It is necessary to define new listers when new routes are defined
        # in the Bottle application, or the static site generation routine
        # will not be able to detect the files to export.
        self.file_listers = [base_lister]

        def file_renderer(path):
            if is_public(path):
                if path.split(".")[-1] in self.template_exts and self.templates.has_template(path):
                    f = File(fname=path, root=self.folder)
                    t = self.templates.get_template(path)
                    if self.refresh:
                        self.files = self.collect_files
                    data = t.render_unicode(p=self, f=f, u=utils)
                    return data.encode(t.module._source_encoding)
                return bottle.static_file(path, root=self.folder)
            return bottle.HTTPError(404, 'File does not exist.')

            # The default function used to render files. Could be modified to change the way files are

        # generated, like using another template language or transforming css...
        self.file_renderer = file_renderer
        self.app.route('/', method=['GET', 'POST', 'PUT', 'DELETE'])(lambda: self.file_renderer('index.html'))
        self.app.route('/<path:path>', method=['GET', 'POST', 'PUT', 'DELETE'])(lambda path: self.file_renderer(path))

    def set_folder(self, folder):
        """
        Sets the folder where the files to serve are located.
        """
        self.folder = folder
        self.templates.directories[0] = folder
        if self.folder not in sys.path:
            sys.path.append(self.folder)

        # append more ignored patterns
        ignore_file = op.join(self.folder, ".ignore")
        if os.path.isfile(ignore_file):
            patts = list(file(ignore_file))
            patts = map(lambda x: x.strip(), patts)
            patts = filter(lambda x: not x.startswith("#"), patts)
            self.file_exclusion.extend(patts)

    def run(self, host='0.0.0.0', port=8080):
        """
        Launch a development web server.
        """
        waitress.serve(self, host=host, port=port)

    def get(self, f):
        """
        Get the content of a file, indentified by its path relative to the folder configured
        in PyGreen. If the file extension is one of the extensions that should be processed
        through Mako, it will be processed.
        """
        handler = wsgiref.handlers.SimpleHandler(sys.stdin, sys.stdout, sys.stderr, {})
        handler.setup_environ()
        env = handler.environ
        env.update({'PATH_INFO': "/%s" % f.fname, 'REQUEST_METHOD': "GET"})
        out = b"".join(self.app(env, lambda *args: None))
        return out

    @property
    def settings(self):
        m = __import__('settings', globals(), locals(), [], -1)
        reload(m)
        return m

    def link(self, fname, start=None):

        items = filter(lambda x: re.search(fname, x.fname, re.IGNORECASE), self.files)
        if not items:
            f = self.files[0]
            _logger.error("*** name %s does not match" % fname)
        else:
            f = items[0]
            if len(items) > 1:
                _logger.warn("name %s matches more than one item" % fname)

        link, name = f.url(start)
        return '%s' % (link, name)

    def toc(self, tag=None, match=None, start=None):
        "Produces name, links pairs from file names"

        if tag:
            items = filter(lambda x: tag in x.meta['tags'], self.files)
        else:
            items = self.files

        if match:
            items = filter(lambda x: re.search(match, x.fname, re.IGNORECASE), self.files)

        if not items:
            _logger.error("*** tag %s does not match" % tag)

        urls = [f.url(start) for f in items]
        t = self.templates.get_template("toc.html")
        h = t.render_unicode(urls=urls)
        return h

    def gallery(self, match=None, start=None, span=4):
        "Produces name, links pairs from file names"

        items = filter(lambda x: x.is_image, self.files)

        if match:
            items = filter(lambda x: re.search(match, x.fname, re.IGNORECASE), self.files)

        urls = [ f.url(start) for f in items ]
        t = self.templates.get_template("gallery.html")
        h = t.render_unicode(urls=urls, span=span)
        return h

    @property
    def collect_files(self):
        """
        Collects all files that will be parsed. Will also be available in main context.
        TODO: this method crawls the entire directory tree each time it is accessed.
        It is handy during development but very large trees may affect performance.
        """
        files = []
        for l in self.file_listers:
            files += l()

        files = map(lambda x: File(x, self.folder), files)

        # apply sort order
        decor = [(f.sortkey, f.name, f) for f in files]
        decor.sort()
        return [ f[2] for f in decor ]

    def gen_static(self, output_folder):
        """
        Generates a complete static version of the web site. It will stored in
        output_folder.
        """

        # this makes all files available in the template context
        self.files = self.collect_files
        self.files = filter(self.is_public, self.files)
        for f in self.files:
            if f.skip_file:
                _logger.info("skipping large file %s of %.1fkb" % (f.fname, f.size))
                continue
            _logger.info("generating %s" % f.fname)
            content = self.get(f)
            f.write(output_folder, content)

    def __call__(self, environ, start_response):
        return self.app(environ, start_response)

    def cli(self, cmd_args=None):
        """
        The command line interface of PyGreen.
        """
        logging.basicConfig(level=logging.INFO, format='%(message)s')

        parser = argparse.ArgumentParser(description='PyBlue, micro web framework/static web site generator')
        subparsers = parser.add_subparsers(dest='action')

        parser_serve = subparsers.add_parser('serve', help='serve the web site')
        parser_serve.add_argument('-p', '--port', type=int, default=8080, help='folder containg files to serve')
        parser_serve.add_argument('-f', '--folder', default=".", help='folder containg files to serve')
        parser_serve.add_argument('-d', '--disable-templates', action='store_true', default=False,
                                  help='just serve static files, do not use invoke Mako')

        def serve():
            if args.disable_templates:
                self.template_exts = set([])
            self.run(port=args.port)

        parser_serve.set_defaults(func=serve)

        parser_gen = subparsers.add_parser('gen', help='generate a static version of the site')
        parser_gen.add_argument('output', help='folder to store the files')
        parser_gen.add_argument('-f', '--folder', default=".", help='folder containg files to serve')

        def gen():
            self.refresh = False
            self.gen_static(args.output)

        parser_gen.set_defaults(func=gen)

        args = parser.parse_args(cmd_args)
        self.set_folder(args.folder)
        print(parser.description)
        print("")
        args.func()

pygreen = PyGreen()

if __name__ == "__main__":
    pygreen.cli()
<commit_msg>bugfix for site gen<commit_after>#! /usr/bin/python

# PyGreen
# Copyright (c) 2013, Nicolas Vanhoren
#
# Released under the MIT license
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
# Software, and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN
# AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

from __future__ import unicode_literals, print_function

import bottle
import os.path
from mako.lookup import TemplateLookup
import os, os.path, itertools
import wsgiref.handlers
import sys, logging, re
import argparse
import sys
import markdown
import waitress
import utils

_logger = logging.getLogger(__name__)

# always adds the location of the default templates

op = os.path
dn = op.dirname

class File(object):
    "Represents a file object within PyBlue"

    def __init__(self, fname, root):
        self.root  = root
        self.fname = fname
        self.fpath = os.path.join(root, fname)
        self.dname = dn(self.fpath)
        self.ext   = os.path.splitext(fname)[1]

        self.meta  =  dict(name=self.nice_name, sortkey="2", tags=set("data"))
        if not self.skip_file:
            self.meta.update(utils.parse_meta(self.fpath))

    @property
    def nice_name(self):
        "Attempts to generate a nicer name from the filename"
        head, tail = os.path.split(self.fname)
        base, ext = os.path.splitext(tail)
        name = base.title().replace("-", " ").replace("_", " ")
        if self.is_image:
            name = name + self.ext
        return name

    def content(self):
        return file(self.fpath).read()

    @property
    def is_image(self):
        return self.ext in (".png", ".jpg", ".gif")

    @property
    def name(self):
        name = self.meta.get("name", "*** name not set ***")
        return name

    @property
    def sortkey(self):
        return self.meta.get("sortkey", "2")

    @property
    def size(self):
        "File size in KB"
        return utils.get_size(self.fpath)

    @property
    def skip_file(self):
        return self.size > utils.MAX_SIZE_KB

    def write(self, output_folder, text):
        loc = os.path.join(output_folder, self.fname)
        if op.abspath(loc) == op.abspath(self.fpath):
            raise Exception("may not overwrite the original file %s" % loc)

        d = os.path.dirname(loc)
        if not os.path.exists(d):
            os.makedirs(d)
        with open(loc, "wb") as fp:
            fp.write(text)

    def url(self, start=None):

        start = start or self

        rpath = op.relpath(self.root, start.dname)
        rpath = op.join(rpath, self.fname)

        return rpath, self.name

    def __repr__(self):
        return "File: %s (%s)" % (self.name, self.fname)

class PyGreen:
    TEMPLATE_DIR = op.abspath(op.join(op.split(__file__)[0], "templates"))

    def __init__(self):

        # the Bottle application
        self.app = bottle.Bottle()
        # a set of strings that identifies the extension of the files
        # that should be processed using Mako
        self.template_exts = set(["html"])
        # the folder where the files to serve are located. Do not set
        # directly, use set_folder instead
        self.folder = "."

        # recrawls directories on each request, use in serving
        self.refresh = True

        # the collection of files
        self.files = None

        # the TemplateLookup of Mako
        self.templates = TemplateLookup(directories=[self.folder, self.TEMPLATE_DIR],
                                        imports=["from markdown import markdown"],
                                        input_encoding='iso-8859-1',
                                        collection_size=100,
        )
        # A list of regular expression. Files whose the name match
        # one of those regular expressions will not be outputed when generating
        # a static version of the web site
        self.file_exclusion = [r".*\.mako", r".*\.py", r"(^|.*\/)\..*"]

        def is_public(path):
            for ex in self.file_exclusion:
                if re.match(ex, path):
                    return False
            return True

        def base_lister():
            files = []
            for dirpath, dirnames, filenames in os.walk(self.folder):
                filenames.sort()
                for f in filenames:
                    absp = os.path.join(dirpath, f)
                    path = os.path.relpath(absp, self.folder)
                    if is_public(path):
                        files.append(path)
            return files

            # A list of function. Each function must return a list of paths

        # of files to export during the generation of the static web site.
        # The default one simply returns all the files contained in the folder.
        # It is necessary to define new listers when new routes are defined
        # in the Bottle application, or the static site generation routine
        # will not be able to detect the files to export.
        self.file_listers = [base_lister]

        def file_renderer(path):
            if is_public(path):
                if path.split(".")[-1] in self.template_exts and self.templates.has_template(path):
                    f = File(fname=path, root=self.folder)
                    t = self.templates.get_template(path)
                    if self.refresh:
                        self.files = self.collect_files
                    data = t.render_unicode(p=self, f=f, u=utils)
                    return data.encode(t.module._source_encoding)
                return bottle.static_file(path, root=self.folder)
            return bottle.HTTPError(404, 'File does not exist.')

            # The default function used to render files. Could be modified to change the way files are

        # generated, like using another template language or transforming css...
        self.file_renderer = file_renderer
        self.app.route('/', method=['GET', 'POST', 'PUT', 'DELETE'])(lambda: self.file_renderer('index.html'))
        self.app.route('/<path:path>', method=['GET', 'POST', 'PUT', 'DELETE'])(lambda path: self.file_renderer(path))

    def set_folder(self, folder):
        """
        Sets the folder where the files to serve are located.
        """
        self.folder = folder
        self.templates.directories[0] = folder
        if self.folder not in sys.path:
            sys.path.append(self.folder)

        # append more ignored patterns
        ignore_file = op.join(self.folder, ".ignore")
        if os.path.isfile(ignore_file):
            patts = list(file(ignore_file))
            patts = map(lambda x: x.strip(), patts)
            patts = filter(lambda x: not x.startswith("#"), patts)
            self.file_exclusion.extend(patts)

    def run(self, host='0.0.0.0', port=8080):
        """
        Launch a development web server.
        """
        waitress.serve(self, host=host, port=port)

    def get(self, f):
        """
        Get the content of a file, indentified by its path relative to the folder configured
        in PyGreen. If the file extension is one of the extensions that should be processed
        through Mako, it will be processed.
        """
        handler = wsgiref.handlers.SimpleHandler(sys.stdin, sys.stdout, sys.stderr, {})
        handler.setup_environ()
        env = handler.environ
        env.update({'PATH_INFO': "/%s" % f.fname, 'REQUEST_METHOD': "GET"})
        out = b"".join(self.app(env, lambda *args: None))
        return out

    @property
    def settings(self):
        m = __import__('settings', globals(), locals(), [], -1)
        reload(m)
        return m

    def link(self, fname, start=None):

        items = filter(lambda x: re.search(fname, x.fname, re.IGNORECASE), self.files)
        if not items:
            f = self.files[0]
            _logger.error("*** name %s does not match" % fname)
        else:
            f = items[0]
            if len(items) > 1:
                _logger.warn("name %s matches more than one item" % fname)

        link, name = f.url(start)
        return '%s' % (link, name)

    def toc(self, tag=None, match=None, start=None):
        "Produces name, links pairs from file names"

        if tag:
            items = filter(lambda x: tag in x.meta['tags'], self.files)
        else:
            items = self.files

        if match:
            items = filter(lambda x: re.search(match, x.fname, re.IGNORECASE), self.files)

        if not items:
            _logger.error("*** tag %s does not match" % tag)

        urls = [f.url(start) for f in items]
        t = self.templates.get_template("toc.html")
        h = t.render_unicode(urls=urls)
        return h

    def gallery(self, match=None, start=None, span=4):
        "Produces name, links pairs from file names"

        items = filter(lambda x: x.is_image, self.files)

        if match:
            items = filter(lambda x: re.search(match, x.fname, re.IGNORECASE), self.files)

        urls = [ f.url(start) for f in items ]
        t = self.templates.get_template("gallery.html")
        h = t.render_unicode(urls=urls, span=span)
        return h

    @property
    def collect_files(self):
        """
        Collects all files that will be parsed. Will also be available in main context.
        TODO: this method crawls the entire directory tree each time it is accessed.
        It is handy during development but very large trees may affect performance.
        """
        files = []
        for l in self.file_listers:
            files += l()

        files = map(lambda x: File(x, self.folder), files)

        # apply sort order
        decor = [(f.sortkey, f.name, f) for f in files]
        decor.sort()
        return [ f[2] for f in decor ]

    def gen_static(self, output_folder):
        """
        Generates a complete static version of the web site. It will stored in
        output_folder.
        """

        # this makes all files available in the template context
        self.files = self.collect_files

        for f in self.files:
            if f.skip_file:
                _logger.info("skipping large file %s of %.1fkb" % (f.fname, f.size))
                continue
            _logger.info("generating %s" % f.fname)
            content = self.get(f)
            f.write(output_folder, content)

    def __call__(self, environ, start_response):
        return self.app(environ, start_response)

    def cli(self, cmd_args=None):
        """
        The command line interface of PyGreen.
        """
        logging.basicConfig(level=logging.INFO, format='%(message)s')

        parser = argparse.ArgumentParser(description='PyBlue, micro web framework/static web site generator')
        subparsers = parser.add_subparsers(dest='action')

        parser_serve = subparsers.add_parser('serve', help='serve the web site')
        parser_serve.add_argument('-p', '--port', type=int, default=8080, help='folder containg files to serve')
        parser_serve.add_argument('-f', '--folder', default=".", help='folder containg files to serve')
        parser_serve.add_argument('-d', '--disable-templates', action='store_true', default=False,
                                  help='just serve static files, do not use invoke Mako')

        def serve():
            if args.disable_templates:
                self.template_exts = set([])
            self.run(port=args.port)

        parser_serve.set_defaults(func=serve)

        parser_gen = subparsers.add_parser('gen', help='generate a static version of the site')
        parser_gen.add_argument('output', help='folder to store the files')
        parser_gen.add_argument('-f', '--folder', default=".", help='folder containg files to serve')

        def gen():
            self.refresh = False
            self.gen_static(args.output)

        parser_gen.set_defaults(func=gen)

        args = parser.parse_args(cmd_args)
        self.set_folder(args.folder)
        print(parser.description)
        print("")
        args.func()

pygreen = PyGreen()

if __name__ == "__main__":
    pygreen.cli()
<|endoftext|>


<commit_before>"""Tests for context.py."""

import os
import re
import sys
import time
import unittest

from google.appengine.api import apiproxy_stub_map
from google.appengine.api import datastore_file_stub
from google.appengine.api import memcache
from google.appengine.api.memcache import memcache_stub

from core import datastore_rpc
from core import datastore_query

from ndb import context
from ndb import eventloop
from ndb import model
from ndb import tasks

class MyAutoBatcher(context.AutoBatcher):

  _log = []

  @classmethod
  def reset_log(cls):
    cls._log = []

  def __init__(self, todo_task):
    def wrap(*args):
      self.__class__._log.append(args)
      return todo_task(*args)
    super(MyAutoBatcher, self).__init__(wrap)

class ContextTests(unittest.TestCase):

  def setUp(self):
    self.set_up_eventloop()
    self.set_up_stubs()
    MyAutoBatcher.reset_log()
    self.ctx = context.Context(auto_batcher_class=MyAutoBatcher)

  def set_up_eventloop(self):
    if eventloop._EVENT_LOOP_KEY in os.environ:
      del os.environ[eventloop._EVENT_LOOP_KEY]
    self.ev = eventloop.get_event_loop()
    self.log = []

  def set_up_stubs(self):
    apiproxy_stub_map.apiproxy = apiproxy_stub_map.APIProxyStubMap()
    ds_stub = datastore_file_stub.DatastoreFileStub('_', None)
    apiproxy_stub_map.apiproxy.RegisterStub('datastore_v3', ds_stub)
    mc_stub = memcache_stub.MemcacheServiceStub()
    apiproxy_stub_map.apiproxy.RegisterStub('memcache', mc_stub)

  def testContext_AutoBatcher_Get(self):
    @tasks.task
    def foo():
      key1 = model.Key(flat=['Foo', 1])
      key2 = model.Key(flat=['Foo', 2])
      key3 = model.Key(flat=['Foo', 3])
      fut1 = self.ctx.get(key1)
      fut2 = self.ctx.get(key2)
      fut3 = self.ctx.get(key3)
      ent1 = yield fut1
      ent2 = yield fut2
      ent3 = yield fut3
      raise tasks.Return([ent1, ent2, ent3])
    ents = foo().get_result()
    self.assertEqual(ents, [None, None, None])
    self.assertEqual(len(MyAutoBatcher._log), 1)

  @tasks.task
  def create_entities(self):
    key0 = model.Key(flat=['Foo', None])
    ent1 = model.Model(key=key0)
    ent2 = model.Model(key=key0)
    ent3 = model.Model(key=key0)
    fut1 = self.ctx.put(ent1)
    fut2 = self.ctx.put(ent2)
    fut3 = self.ctx.put(ent3)
    key1 = yield fut1
    key2 = yield fut2
    key3 = yield fut3
    raise tasks.Return([key1, key2, key3])

  def testContext_AutoBatcher_Put(self):
    keys = self.create_entities().get_result()
    self.assertEqual(len(keys), 3)
    self.assertTrue(None not in keys)
    self.assertEqual(len(MyAutoBatcher._log), 1)

  def testContext_AutoBatcher_Delete(self):
    @tasks.task
    def foo():
      key1 = model.Key(flat=['Foo', 1])
      key2 = model.Key(flat=['Foo', 2])
      key3 = model.Key(flat=['Foo', 3])
      fut1 = self.ctx.delete(key1)
      fut2 = self.ctx.delete(key2)
      fut3 = self.ctx.delete(key3)
      yield fut1
      yield fut2
      yield fut3
    foo().check_success()
    self.assertEqual(len(MyAutoBatcher._log), 1)

  def testContext_Cache(self):
    @tasks.task
    def foo():
      key1 = model.Key(flat=('Foo', 1))
      ent1 = model.Expando(key=key1, foo=42, bar='hello')
      key = yield self.ctx.put(ent1)
      self.assertTrue(key1 in self.ctx._cache)  # Whitebox.
      a = yield self.ctx.get(key1)
      b = yield self.ctx.get(key1)
      self.assertTrue(a is b)
      yield self.ctx.delete(key1)
      self.assertTrue(self.ctx._cache[key] is None)  # Whitebox.
      a = yield self.ctx.get(key1)
      self.assertTrue(a is None)
    foo().check_success()

  def testContext_CachePolicy(self):
    def should_cache(key):
      return False
    @tasks.task
    def foo():
      key1 = model.Key(flat=('Foo', 1))
      ent1 = model.Expando(key=key1, foo=42, bar='hello')
      key = yield self.ctx.put(ent1)
      self.assertTrue(key1 not in self.ctx._cache)  # Whitebox.
      a = yield self.ctx.get(key1)
      b = yield self.ctx.get(key1)
      self.assertTrue(a is not b)
      yield self.ctx.delete(key1)
      self.assertTrue(key not in self.ctx._cache)  # Whitebox.
      a = yield self.ctx.get(key1)
      self.assertTrue(a is None)
    self.ctx.set_cache_policy(should_cache)
    foo().check_success()

  def testContext_Memcache(self):
    @tasks.task
    def foo():
      key1 = model.Key(flat=('Foo', 1))
      key2 = model.Key(flat=('Foo', 2))
      ent1 = model.Expando(key=key1, foo=42, bar='hello')
      ent2 = model.Expando(key=key2, foo=1, bar='world')
      k1, k2 = yield self.ctx.put(ent1), self.ctx.put(ent2)
      self.assertEqual(k1, key1)
      self.assertEqual(k2, key2)
      yield tasks.sleep(0.01)  # Let other task complete.
      keys = [k1.urlsafe(), k2.urlsafe()]
      results = memcache.get_multi(keys)
      self.assertEqual(
        results,
        {key1.urlsafe(): self.ctx._conn.adapter.entity_to_pb(ent1),
         key2.urlsafe(): self.ctx._conn.adapter.entity_to_pb(ent2)})
    foo().check_success()

  def testContext_CacheQuery(self):
    @tasks.task
    def foo():
      key1 = model.Key(flat=('Foo', 1))
      key2 = model.Key(flat=('Foo', 2))
      ent1 = model.Expando(key=key1, foo=42, bar='hello')
      ent2 = model.Expando(key=key2, foo=1, bar='world')
      key1a, key2a = yield self.ctx.put(ent1),  self.ctx.put(ent2)
      self.assertTrue(key1 in self.ctx._cache)  # Whitebox.
      self.assertTrue(key2 in self.ctx._cache)  # Whitebox.
      self.assertEqual(key1, key1a)
      self.assertEqual(key2, key2a)
      @tasks.task
      def callback(ent):
        return ent
      query = datastore_query.Query(app='_', kind='Foo')
      results = yield self.ctx.map_query(query, callback)
      self.assertEqual(results, [ent1, ent2])
      self.assertTrue(results[0] is ent1)
      self.assertTrue(results[1] is ent2)
    foo().check_success()

  def testContext_AllocateIds(self):
    @tasks.task
    def foo():
      key = model.Key(flat=('Foo', 1))
      lo_hi = yield self.ctx.allocate_ids(key, size=10)
      self.assertEqual(lo_hi, (1, 10))
      lo_hi = yield self.ctx.allocate_ids(key, max=20)
      self.assertEqual(lo_hi, (11, 20))
    foo().check_success()

  def testContext_MapQuery(self):
    @tasks.task
    def callback(ent):
      return ent.key.flat()[-1]
    @tasks.task
    def foo():
      yield self.create_entities()
      query = datastore_query.Query(app='_', kind='Foo')
      res = yield self.ctx.map_query(query, callback)
      raise tasks.Return(res)
    res = foo().get_result()
    self.assertEqual(set(res), set([1, 2, 3]))

  def testContext_MapQuery_NoCallback(self):
    @tasks.task
    def foo():
      yield self.create_entities()
      query = datastore_query.Query(app='_', kind='Foo')
      res = yield self.ctx.map_query(query, None)
      raise tasks.Return(res)
    res = foo().get_result()
    self.assertEqual(len(res), 3)
    for i, ent in enumerate(res):
      self.assertTrue(isinstance(ent, model.Model))
      self.assertEqual(ent.key.flat(), ['Foo', i+1])

  def testContext_MapQuery_NonTaskCallback(self):
    def callback(ent):
      return ent.key.flat()[-1]
    @tasks.task
    def foo():
      yield self.create_entities()
      query = datastore_query.Query(app='_', kind='Foo')
      res = yield self.ctx.map_query(query, callback)
      raise tasks.Return(res)
    res = foo().get_result()
    self.assertEqual(res, [1, 2, 3])

  def testContext_TransactionFailed(self):
    @tasks.task
    def foo():
      key = model.Key(flat=('Foo', 1))
      ent = model.Expando(key=key, bar=1)
      yield self.ctx.put(ent)
      @tasks.task
      def callback(ctx):
        self.assertTrue(key not in ctx._cache)  # Whitebox.
        e = yield ctx.get(key)
        self.assertTrue(key in ctx._cache)  # Whitebox.
        e.bar = 2
        yield ctx.put(e)
      yield self.ctx.transaction(callback)
      self.assertEqual(self.ctx._cache[key].bar, 2)
    foo().check_success()

  def testContext_GetOrInsert(self):
    # This also tests Context.transaction()
    class Mod(model.Model):
      data = model.StringProperty()
    @tasks.task
    def foo():
      ent = yield self.ctx.get_or_insert(Mod, 'a', data='hello')
      assert isinstance(ent, Mod)
      ent2 = yield self.ctx.get_or_insert(Mod, 'a', data='hello')
      assert ent2 == ent
    foo().check_success()

  def testAddContextDecorator(self):
    class Demo(object):
      @context.toplevel
      def method(self, arg):
        return (self.ctx, arg)
    a = Demo()
    self.assertFalse(hasattr(a, 'ctx'))
    ctx, arg = a.method(42)
    self.assertTrue(isinstance(ctx, context.Context))
    self.assertEqual(arg, 42)

def main():
  unittest.main()

if __name__ == '__main__':
  main()
<commit_msg>Add unittest for default context in transaction.<commit_after>"""Tests for context.py."""

import os
import re
import sys
import time
import unittest

from google.appengine.api import apiproxy_stub_map
from google.appengine.api import datastore_file_stub
from google.appengine.api import memcache
from google.appengine.api.memcache import memcache_stub

from core import datastore_rpc
from core import datastore_query

from ndb import context
from ndb import eventloop
from ndb import model
from ndb import tasks

class MyAutoBatcher(context.AutoBatcher):

  _log = []

  @classmethod
  def reset_log(cls):
    cls._log = []

  def __init__(self, todo_task):
    def wrap(*args):
      self.__class__._log.append(args)
      return todo_task(*args)
    super(MyAutoBatcher, self).__init__(wrap)

class ContextTests(unittest.TestCase):

  def setUp(self):
    self.set_up_eventloop()
    self.set_up_stubs()
    MyAutoBatcher.reset_log()
    self.ctx = context.Context(auto_batcher_class=MyAutoBatcher)

  def set_up_eventloop(self):
    if eventloop._EVENT_LOOP_KEY in os.environ:
      del os.environ[eventloop._EVENT_LOOP_KEY]
    self.ev = eventloop.get_event_loop()
    self.log = []

  def set_up_stubs(self):
    apiproxy_stub_map.apiproxy = apiproxy_stub_map.APIProxyStubMap()
    ds_stub = datastore_file_stub.DatastoreFileStub('_', None)
    apiproxy_stub_map.apiproxy.RegisterStub('datastore_v3', ds_stub)
    mc_stub = memcache_stub.MemcacheServiceStub()
    apiproxy_stub_map.apiproxy.RegisterStub('memcache', mc_stub)

  def testContext_AutoBatcher_Get(self):
    @tasks.task
    def foo():
      key1 = model.Key(flat=['Foo', 1])
      key2 = model.Key(flat=['Foo', 2])
      key3 = model.Key(flat=['Foo', 3])
      fut1 = self.ctx.get(key1)
      fut2 = self.ctx.get(key2)
      fut3 = self.ctx.get(key3)
      ent1 = yield fut1
      ent2 = yield fut2
      ent3 = yield fut3
      raise tasks.Return([ent1, ent2, ent3])
    ents = foo().get_result()
    self.assertEqual(ents, [None, None, None])
    self.assertEqual(len(MyAutoBatcher._log), 1)

  @tasks.task
  def create_entities(self):
    key0 = model.Key(flat=['Foo', None])
    ent1 = model.Model(key=key0)
    ent2 = model.Model(key=key0)
    ent3 = model.Model(key=key0)
    fut1 = self.ctx.put(ent1)
    fut2 = self.ctx.put(ent2)
    fut3 = self.ctx.put(ent3)
    key1 = yield fut1
    key2 = yield fut2
    key3 = yield fut3
    raise tasks.Return([key1, key2, key3])

  def testContext_AutoBatcher_Put(self):
    keys = self.create_entities().get_result()
    self.assertEqual(len(keys), 3)
    self.assertTrue(None not in keys)
    self.assertEqual(len(MyAutoBatcher._log), 1)

  def testContext_AutoBatcher_Delete(self):
    @tasks.task
    def foo():
      key1 = model.Key(flat=['Foo', 1])
      key2 = model.Key(flat=['Foo', 2])
      key3 = model.Key(flat=['Foo', 3])
      fut1 = self.ctx.delete(key1)
      fut2 = self.ctx.delete(key2)
      fut3 = self.ctx.delete(key3)
      yield fut1
      yield fut2
      yield fut3
    foo().check_success()
    self.assertEqual(len(MyAutoBatcher._log), 1)

  def testContext_Cache(self):
    @tasks.task
    def foo():
      key1 = model.Key(flat=('Foo', 1))
      ent1 = model.Expando(key=key1, foo=42, bar='hello')
      key = yield self.ctx.put(ent1)
      self.assertTrue(key1 in self.ctx._cache)  # Whitebox.
      a = yield self.ctx.get(key1)
      b = yield self.ctx.get(key1)
      self.assertTrue(a is b)
      yield self.ctx.delete(key1)
      self.assertTrue(self.ctx._cache[key] is None)  # Whitebox.
      a = yield self.ctx.get(key1)
      self.assertTrue(a is None)
    foo().check_success()

  def testContext_CachePolicy(self):
    def should_cache(key):
      return False
    @tasks.task
    def foo():
      key1 = model.Key(flat=('Foo', 1))
      ent1 = model.Expando(key=key1, foo=42, bar='hello')
      key = yield self.ctx.put(ent1)
      self.assertTrue(key1 not in self.ctx._cache)  # Whitebox.
      a = yield self.ctx.get(key1)
      b = yield self.ctx.get(key1)
      self.assertTrue(a is not b)
      yield self.ctx.delete(key1)
      self.assertTrue(key not in self.ctx._cache)  # Whitebox.
      a = yield self.ctx.get(key1)
      self.assertTrue(a is None)
    self.ctx.set_cache_policy(should_cache)
    foo().check_success()

  def testContext_Memcache(self):
    @tasks.task
    def foo():
      key1 = model.Key(flat=('Foo', 1))
      key2 = model.Key(flat=('Foo', 2))
      ent1 = model.Expando(key=key1, foo=42, bar='hello')
      ent2 = model.Expando(key=key2, foo=1, bar='world')
      k1, k2 = yield self.ctx.put(ent1), self.ctx.put(ent2)
      self.assertEqual(k1, key1)
      self.assertEqual(k2, key2)
      yield tasks.sleep(0.01)  # Let other task complete.
      keys = [k1.urlsafe(), k2.urlsafe()]
      results = memcache.get_multi(keys)
      self.assertEqual(
        results,
        {key1.urlsafe(): self.ctx._conn.adapter.entity_to_pb(ent1),
         key2.urlsafe(): self.ctx._conn.adapter.entity_to_pb(ent2)})
    foo().check_success()

  def testContext_CacheQuery(self):
    @tasks.task
    def foo():
      key1 = model.Key(flat=('Foo', 1))
      key2 = model.Key(flat=('Foo', 2))
      ent1 = model.Expando(key=key1, foo=42, bar='hello')
      ent2 = model.Expando(key=key2, foo=1, bar='world')
      key1a, key2a = yield self.ctx.put(ent1),  self.ctx.put(ent2)
      self.assertTrue(key1 in self.ctx._cache)  # Whitebox.
      self.assertTrue(key2 in self.ctx._cache)  # Whitebox.
      self.assertEqual(key1, key1a)
      self.assertEqual(key2, key2a)
      @tasks.task
      def callback(ent):
        return ent
      query = datastore_query.Query(app='_', kind='Foo')
      results = yield self.ctx.map_query(query, callback)
      self.assertEqual(results, [ent1, ent2])
      self.assertTrue(results[0] is ent1)
      self.assertTrue(results[1] is ent2)
    foo().check_success()

  def testContext_AllocateIds(self):
    @tasks.task
    def foo():
      key = model.Key(flat=('Foo', 1))
      lo_hi = yield self.ctx.allocate_ids(key, size=10)
      self.assertEqual(lo_hi, (1, 10))
      lo_hi = yield self.ctx.allocate_ids(key, max=20)
      self.assertEqual(lo_hi, (11, 20))
    foo().check_success()

  def testContext_MapQuery(self):
    @tasks.task
    def callback(ent):
      return ent.key.flat()[-1]
    @tasks.task
    def foo():
      yield self.create_entities()
      query = datastore_query.Query(app='_', kind='Foo')
      res = yield self.ctx.map_query(query, callback)
      raise tasks.Return(res)
    res = foo().get_result()
    self.assertEqual(set(res), set([1, 2, 3]))

  def testContext_MapQuery_NoCallback(self):
    @tasks.task
    def foo():
      yield self.create_entities()
      query = datastore_query.Query(app='_', kind='Foo')
      res = yield self.ctx.map_query(query, None)
      raise tasks.Return(res)
    res = foo().get_result()
    self.assertEqual(len(res), 3)
    for i, ent in enumerate(res):
      self.assertTrue(isinstance(ent, model.Model))
      self.assertEqual(ent.key.flat(), ['Foo', i+1])

  def testContext_MapQuery_NonTaskCallback(self):
    def callback(ent):
      return ent.key.flat()[-1]
    @tasks.task
    def foo():
      yield self.create_entities()
      query = datastore_query.Query(app='_', kind='Foo')
      res = yield self.ctx.map_query(query, callback)
      raise tasks.Return(res)
    res = foo().get_result()
    self.assertEqual(res, [1, 2, 3])

  def testContext_TransactionFailed(self):
    @tasks.task
    def foo():
      key = model.Key(flat=('Foo', 1))
      ent = model.Expando(key=key, bar=1)
      yield self.ctx.put(ent)
      @tasks.task
      def callback(ctx):
        self.assertTrue(key not in ctx._cache)  # Whitebox.
        e = yield ctx.get(key)
        self.assertTrue(key in ctx._cache)  # Whitebox.
        e.bar = 2
        yield ctx.put(e)
      yield self.ctx.transaction(callback)
      self.assertEqual(self.ctx._cache[key].bar, 2)
    foo().check_success()

  def testContext_GetOrInsert(self):
    # This also tests Context.transaction()
    class Mod(model.Model):
      data = model.StringProperty()
    @tasks.task
    def foo():
      ent = yield self.ctx.get_or_insert(Mod, 'a', data='hello')
      assert isinstance(ent, Mod)
      ent2 = yield self.ctx.get_or_insert(Mod, 'a', data='hello')
      assert ent2 == ent
    foo().check_success()

  def testAddContextDecorator(self):
    class Demo(object):
      @context.toplevel
      def method(self, arg):
        return (self.ctx, arg)
    a = Demo()
    self.assertFalse(hasattr(a, 'ctx'))
    ctx, arg = a.method(42)
    self.assertTrue(isinstance(ctx, context.Context))
    self.assertEqual(arg, 42)

  def testDefaultContextTransaction(self):
    @context.taskify
    def outer():
      ctx1 = context._get_default_context()
      @context.task
      def inner():
        ctx2 = context._get_default_context()
        self.assertTrue(ctx1 is not ctx2)
        self.assertTrue(isinstance(ctx2._conn,
                                   datastore_rpc.TransactionalConnection))
        return 42
      a = yield context.transaction(inner)
      ctx1a = context._get_default_context()
      self.assertTrue(ctx1 is ctx1a)
      raise tasks.Return(a)
    b = outer()
    self.assertEqual(b, 42)

def main():
  unittest.main()

if __name__ == '__main__':
  main()
<|endoftext|>


<commit_before>from __future__ import (absolute_import, division, print_function,
        unicode_literals)
from builtins import (ascii, bytes, chr, dict, filter, hex, input, str, super,
        zip)
from future.standard_library import install_aliases  # for urlopen
install_aliases()

import unittest
from urllib.request import urlopen

from selenium import webdriver
from selenium.webdriver.common.keys import Keys

from .base import SeleniumTest

class GameTest(SeleniumTest):

    def count_stones_and_points(self):
        imgs = self.browser.find_elements_by_tag_name('img')
        empty = 0
        black = 0
        white = 0
        for img in imgs:
            if 'e.gif' in img.get_attribute('src'):
                empty += 1
            elif 'b.gif' in img.get_attribute('src'):
                black += 1
            elif 'w.gif' in img.get_attribute('src'):
                white += 1
        return {'empty': empty, 'black': black, 'white': white}

    def find_empty_point_to_click(self):
        links = self.browser.find_elements_by_css_selector('table.goban a')
        target_link = None
        for link in links:
            if ('e.gif' in
                    link.find_element_by_tag_name('img').get_attribute('src')):
                target_link = link
                break
        return target_link

    def test_game_page(self):
        self.browser.get(self.get_server_url() + '/game')
        ## just to make sure page is loaded
        self.wait_for(lambda: self.assertIn('Go', self.browser.title))

        # on the game page is a table with class 'goban'
        table = self.browser.find_element_by_css_selector('table.goban')
            ## should not raise

        # on the game page are 19x19 imgs representing board points/stones
        empty = self.count_stones_and_points()['empty']
        self.assertEqual(19*19, empty, "did not find 19x19 board imgs")

        ## check one of those images can be loaded
        img = self.browser.find_element_by_css_selector('table.goban a img')
        response = urlopen(img.get_attribute('src'))
        self.assertEqual(response.getcode(), 200)
        try:
            self.assertNotIn(
                'Exception', response.read().decode(),
                'image load returns exception'
            )
        except UnicodeDecodeError:
            pass  ## fine, we got image data

        # user clicks an empty spot, which is a link
        try:
            self.find_empty_point_to_click().click()
        except AttributeError:
            self.fail('no clickable board point found')

        # now on the board is one black stone and 19x19 - 1 empty points
        counts = self.count_stones_and_points()
        self.assertEqual(counts['empty'], 19*19-1)
        self.assertEqual(counts['black'], 1)

        # user clicks another empty spot
        try:
            self.find_empty_point_to_click().click()
        except AttributeError:
            self.fail('no clickable board point found')

        # now one black stone, one white, rest empty
        counts = self.count_stones_and_points()
        self.assertEqual(counts['empty'], 19*19-2)
        self.assertEqual(counts['black'], 1)
        self.assertEqual(counts['white'], 1)

if __name__ == '__main__':
    unittest.main()
<commit_msg>messing with namedtuples some more<commit_after>from __future__ import (absolute_import, division, print_function,
        unicode_literals)
from builtins import (ascii, bytes, chr, dict, filter, hex, input, str, super,
        zip)
from future.standard_library import install_aliases  # for urlopen
install_aliases()

from collections import namedtuple
import unittest
from urllib.request import urlopen

from selenium import webdriver
from selenium.webdriver.common.keys import Keys

from .base import SeleniumTest

class GameTest(SeleniumTest):

    Count = namedtuple('Count', ['white', 'black', 'empty'])
    def count_stones_and_points(self):
        imgs = self.browser.find_elements_by_tag_name('img')
        empty = 0
        black = 0
        white = 0
        for img in imgs:
            if 'e.gif' in img.get_attribute('src'):
                empty += 1
            elif 'b.gif' in img.get_attribute('src'):
                black += 1
            elif 'w.gif' in img.get_attribute('src'):
                white += 1
        return GameTest.Count(empty=empty, black=black, white=white)

    def find_empty_point_to_click(self):
        links = self.browser.find_elements_by_css_selector('table.goban a')
        target_link = None
        for link in links:
            if ('e.gif' in
                    link.find_element_by_tag_name('img').get_attribute('src')):
                target_link = link
                break
        return target_link

    def test_game_page(self):
        self.browser.get(self.get_server_url() + '/game')
        ## just to make sure page is loaded
        self.wait_for(lambda: self.assertIn('Go', self.browser.title))

        # on the game page is a table with class 'goban'
        table = self.browser.find_element_by_css_selector('table.goban')
            ## should not raise

        # on the game page are 19x19 imgs representing board points/stones
        empty = self.count_stones_and_points().empty
        self.assertEqual(19*19, empty, "did not find 19x19 board imgs")

        ## check one of those images can be loaded
        img = self.browser.find_element_by_css_selector('table.goban a img')
        response = urlopen(img.get_attribute('src'))
        self.assertEqual(response.getcode(), 200)
        try:
            self.assertNotIn(
                'Exception', response.read().decode(),
                'image load returns exception'
            )
        except UnicodeDecodeError:
            pass  ## fine, we got image data

        # user clicks an empty spot, which is a link
        try:
            self.find_empty_point_to_click().click()
        except AttributeError:
            self.fail('no clickable board point found')

        # now on the board is one black stone and 19x19 - 1 empty points
        counts = self.count_stones_and_points()
        self.assertEqual(counts.empty, 19*19-1)
        self.assertEqual(counts.black, 1)

        # user clicks another empty spot
        try:
            self.find_empty_point_to_click().click()
        except AttributeError:
            self.fail('no clickable board point found')

        # now one black stone, one white, rest empty
        counts = self.count_stones_and_points()
        self.assertEqual(counts.empty, 19*19-2)
        self.assertEqual(counts.black, 1)
        self.assertEqual(counts.white, 1)

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before>from neb.api import TrinityResource

class Relationship(TrinityResource):
    def create(self, from_node, to, link_type, append=None, increment=None,
            **kwargs):
        params = dict(from_node=from_node, to=to, link_type=link_type,
                append=append, increment=increment, data=kwargs)
        return self.post(self._relationship_path(from_node), payload=params)

    @staticmethod
    def _relationship_path(node_id, relationship_id=None):
        if relationship_id:
            path = 'node/%s/relationships/%d' % (node_id, relationship_id)
        else:
            path = 'node/%s/relationships' % node_id
        return path

    def request(self, *args, **kwargs):
        response = super(Relationship, self).request(*args, **kwargs)
        return Relationship(data=response)

    def __str__(self):
        if hasattr(self, 'from_node') and hasattr(self, 'to'):
            return "<%s => %s>" % (self.from_node, self.to)
        else:
            return super(Relationship, self).__str__()
<commit_msg>Only send append/increment if they are specified.<commit_after>from neb.api import TrinityResource

class Relationship(TrinityResource):
    def create(self, from_node, to, link_type, append=None, increment=None,
            **kwargs):
        params = dict(from_node=from_node, to=to, link_type=link_type,
                data=kwargs)
        if append:
            params[append] = append
        if increment:
            params['increment'] = increment
        return self.post(self._relationship_path(from_node), payload=params)

    @staticmethod
    def _relationship_path(node_id, relationship_id=None):
        if relationship_id:
            path = 'node/%s/relationships/%d' % (node_id, relationship_id)
        else:
            path = 'node/%s/relationships' % node_id
        return path

    def request(self, *args, **kwargs):
        response = super(Relationship, self).request(*args, **kwargs)
        return Relationship(data=response)

    def __str__(self):
        if hasattr(self, 'from_node') and hasattr(self, 'to'):
            return "<%s => %s>" % (self.from_node, self.to)
        else:
            return super(Relationship, self).__str__()
<|endoftext|>