<commit_before>import logging

from django.contrib.auth.models import User

from django.core.exceptions import ObjectDoesNotExist

from rest_framework.parsers import JSONParser
from rest_framework.renderers import JSONRenderer
from io import BytesIO

from login.models import Client, ClientLog, ClientStatus
from .client import BaseClient, MQTTException

from ambulance.models import Ambulance
from ambulance.models import Call
from ambulance.serializers import AmbulanceSerializer, AmbulanceUpdateSerializer
from ambulance.serializers import CallSerializer

from hospital.models import Hospital, HospitalEquipment
from hospital.serializers import HospitalSerializer, \
    HospitalEquipmentSerializer

logger = logging.getLogger(__name__)

# SubscribeClient

class SubscribeClient(BaseClient):

    # The callback for when the client receives a CONNACK
    # response from the server.
    def on_connect(self, client, userdata, flags, rc):

        # is connected?
        if not super().on_connect(client, userdata, flags, rc):
            return False

        # Subscribing in on_connect() means that if we lose the
        # connection and reconnect then subscriptions will be renewed.
        # client.subscribe('#', 2)

        # ambulance handler
        self.client.message_callback_add('user/+/ambulance/+/data',
                                         self.on_ambulance)

        # hospital handler
        self.client.message_callback_add('user/+/hospital/+/data',
                                         self.on_hospital)

        # hospital equipment handler
        self.client.message_callback_add('user/+/hospital/+/equipment/+/data',
                                         self.on_hospital_equipment)

        # call handler
        # self.client.message_callback_add('ambulance/+/call',
        #                                 self.on_call)

        # client status handler
        self.client.message_callback_add('user/+/client/+/status',
                                         self.on_client_status)

        # subscribe
        self.subscribe('user/+/ambulance/+/data', 2)
        self.subscribe('user/+/hospital/+/data', 2)
        self.subscribe('user/+/hospital/+/equipment/+/data', 2)
        self.subscribe('user/+/client/+/status', 2)

        if self.verbosity > 0:
            self.stdout.write(self.style.SUCCESS(">> Listening to MQTT messages..."))

        return True

    def send_error_message(self, username, topic, payload, error):

        logger.debug("send_error_message: {}, '{}:{}': '{}'".format(username,
                                                                    topic,
                                                                    payload,
                                                                    error))

        if self.verbosity > 0:
            self.stdout.write(self.style.ERROR("*> Error {}, '{}:{}': {}".format(username,
                                                                                topic,
                                                                                payload,
                                                                                error)))

        try:

            message = JSONRenderer().render({
                'topic': topic,
                'payload': payload,
                'error': error
            })
            self.publish('user/{}/error'.format(username), message)

        except Exception as e:

            logger.warning(('mqtt.SubscribeClient: {}, ' +
                            "topic = '{}:{}', " +
                            "error = '{}', " +
                            "exception = {}").format(username,
                                                     topic,
                                                     payload,
                                                     error,
                                                     e))

    def parse_topic(self, msg):

        if not msg.payload:
            # empty payload
            return

        if self.verbosity > 0:
            self.stdout.write(self.style.SUCCESS(" > Parsing message '{}:{}'".format(msg.topic,
                                                                                msg.payload)))

        # parse topic
        values = msg.topic.split('/')

        # print('values = {}'.format(values))

        try:

            # retrieve user
            # print(User.objects.all())
            user = User.objects.get(username=values[1])

        except ObjectDoesNotExist as e:

            # does not know username
            # cannot send error message to user
            logger.warning(('mqtt.SubscribeClient: {}, ' +
                            "topic = '{}:{}', " +
                            "error = '{}', " +
                            "exception = {}").format(username,
                                                     topic,
                                                     payload,
                                                     error,
                                                     e))
            return

        # parse data
        try:

            # Parse data into json dict
            data = JSONParser().parse(BytesIO(msg.payload))

        except Exception as e:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "JSON formatted incorrectly")
            return

        if len(values) == 5:

            return user, data, values[3]

        elif len(values) == 7:

            return user, data, values[3], values[5]

        else:

            # send error message to user
            # this should never happen because no subscriptions will match
            # topics with different sizes
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Invalid topic")
            return

    # Update ambulance

    def on_ambulance(self, client, userdata, msg):

        logger.debug("on_ambulance: msg = '{}:{}'".format(msg.topic, msg.payload))

        # parse topic
        values = self.parse_topic(msg)
        if not values:
            return

        logger.debug("on_ambulance: values = '{}'".format(values))

        # retrieve parsed values
        user, data, ambulance_id = values

        try:

            logger.debug('ambulance_id = {}'.format(ambulance_id))

            # retrieve ambulance
            ambulance = Ambulance.objects.get(id=ambulance_id)

        except ObjectDoesNotExist:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Ambulance with id '{}' does not exist".format(ambulance_id))
            return

        except Exception as e:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Exception: '{}'".format(e))
            return

        logger.debug('on_ambulance: ambulance = {}'.format(ambulance))

        try:

            is_valid = False
            if isinstance(data, (list, tuple)):

                # update ambulances in bulk
                serializer = AmbulanceUpdateSerializer(data=data,
                                                       many=True,
                                                       partial=True)

                if serializer.is_valid():

                    # save to database
                    serializer.save(ambulance=ambulance, updated_by=user)
                    is_valid = True

            else:

                # update ambulance
                serializer = AmbulanceSerializer(ambulance,
                                                 data=data,
                                                 partial=True)

                if serializer.is_valid():

                    # save to database
                    serializer.save(updated_by=user)
                    is_valid = True

            if not is_valid:

                logger.debug('on_ambulance: INVALID serializer')

                # send error message to user
                self.send_error_message(user, msg.topic, msg.payload,
                                        serializer.errors)

        except Exception as e:

            logger.debug('on_ambulance: serializer EXCEPTION')

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload, e)

        logger.debug('on_ambulance: DONE')

    # Update hospital
    def on_hospital(self, client, userdata, msg):

        logger.debug("on_hospital: msg = '{}:{}'".format(msg.topic, msg.payload))

        # parse topic
        values = self.parse_topic(msg)
        if not values:
            return

        logger.debug("on_hospital: values = '{}'".format(values))

        # retrieve parsed values
        user, data, hospital_id = values

        try:

            # retrieve hospital
            hospital = Hospital.objects.get(id=hospital_id)

        except ObjectDoesNotExist:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Hospital with id '{}' does not exist".format(hospital_id))
            return

        except Exception as e:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Exception: '{}'".format(e))
            return

        logger.debug('on_hospital: hospital = {}'.format(hospital))

        try:

            # update hospital
            serializer = HospitalSerializer(hospital,
                                            data=data,
                                            partial=True)
            if serializer.is_valid():

                logger.debug('on_hospital: valid serializer')

                # save to database
                serializer.save(updated_by=user)

            else:

                logger.debug('on_hospital: INVALID serializer')

                # send error message to user
                self.send_error_message(user, msg.topic, msg.payload,
                                        serializer.errors)

        except Exception as e:

            logger.debug('on_hospital: serializer EXCEPTION')

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload, e)

        logger.debug('on_hospital: DONE')

    # Update hospital equipment
    def on_hospital_equipment(self, client, userdata, msg):

        logger.debug("on_hospital_equipment: msg = '{}:{}'".format(msg.topic, msg.payload))

        # parse topic
        values = self.parse_topic(msg)
        if not values:
            return

        logger.debug("on_hospital_equipment: values = '{}'".format(values))

        # retrieve parsed values
        user, data, hospital_id, equipment_name = values

        try:

            # retrieve hospital equipment
            hospital_equipment = HospitalEquipment.objects.get(hospital=hospital_id,
                                                               equipment__name=equipment_name)

        except ObjectDoesNotExist:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Hospital equipment with hospital id '{}' and name '{}' does not exist".format(
                                        hospital_id, equipment_name))
            return

        except Exception as e:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Exception: '{}'".format(e))
            return

        logger.debug('on_hospital_equipment: equipment = {}'.format(hospital_equipment))

        try:

            # update hospital equipment
            serializer = HospitalEquipmentSerializer(hospital_equipment,
                                                     data=data,
                                                     partial=True)
            if serializer.is_valid():

                logger.debug('on_hospital_equipment: valid serializer')

                # save to database
                serializer.save(updated_by=user)

            else:

                logger.debug('on_hospital_equipment: INVALID serializer')

                # send error message to user
                self.send_error_message(user, msg.topic, msg.payload,
                                        serializer.errors)

        except Exception as e:

            logger.debug('on_hospital_equipment: serializer EXCEPTION')

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload, e)

        logger.debug('on_hospital_equipment: DONE')

    # update client information
    def on_client_status(self, clnt, userdata, msg):

        logger.debug("on_client_status: msg = '{}:{}'".format(msg.topic, msg.payload))

        # parse topic
        values = self.parse_topic(msg)
        if not values:
            return

        logger.debug("on_client_status: values = '{}'".format(values))

        # retrieve parsed values
        user, status, client_id = values

        # handle status
        if status == 'online':

            # user just logged in, create record
            client = Client(client_id=client_id, user=user, status=status)
            client.save()

            # log operation
            log = ClientLog(client=client, action=status)
            log.save()

        elif status == 'offline' or status == 'disconnected':

            # user client offline or disconnected
            try:

                # retrieve client record first
                client = Client.objects.get(client_id=client_id)

                # is online?
                if client.status != ClientStatus.O.name:

                    # client is not online
                    logger.debug('on_client_status: not online')

                    # send error message to user
                    self.send_error_message(user, msg.topic, msg.payload,
                                            "client '{}' is not online".format(client_id))

                    return

                # update status
                client.status = status
                client.save()

                # log operation
                log = ClientLog(client=client, action=status)
                log.save()

                # clean up mqtt topic
                self.remove_topic('/user/{}/client/{}/status'.format(user.username, client_id))

            except Client.DoesNotExist:

                logger.debug('on_client_status: INVALID client')

                # send error message to user
                self.send_error_message(user, msg.topic, msg.payload,
                                        "client '{}' is not valid".format(client_id))

    # update calls
    def on_call(self, client, userdata, msg):

        logger.debug("on_call: msg = '{}:{}'".format(msg.topic, msg.payload))

        # parse topic
        values = self.parse_topic(msg)
        if not values:
            return

        logger.debug("on_call: values = '{}'".format(values))

        # retrieve parsed values
        user, data, call_id = values

        try:

            logger.debug('call_id = {}'.format(call_id))

            call = Call.objects.get(id=call_id)

        except ObjectDoesNotExist:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Call with id '{}' does not exist".format(call_id))

            return

        except Exception as e:
            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Exception: '{}'".format(e))
            return

        logger.debug('on_call: call = {}'.format(call))

        try:

            # update call
            serializer = CallSerializer(call)  # TODO: Find structure

            if serializer.is_valid():
                logger.debug('on_call: vaild serializer')

                # save to database
                serializer.save(updated_by=user)

            else:

                logger.debug('on_call: INVALID serializer')

                # send error message to user
                self.send_error_message(user, msg.topic, msg.payload,
                                        serializer.errors)

        except Exception as e:

            logger.debug('on_call: serializer EXCEPTION')

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload, e)

        logger.debug('on_call: DONE')
<commit_msg>Added Client and ClientLog models to keep track of clients.<commit_after>import logging

from django.contrib.auth.models import User

from django.core.exceptions import ObjectDoesNotExist

from rest_framework.parsers import JSONParser
from rest_framework.renderers import JSONRenderer
from io import BytesIO

from login.models import Client, ClientLog, ClientStatus
from .client import BaseClient, MQTTException

from ambulance.models import Ambulance
from ambulance.models import Call
from ambulance.serializers import AmbulanceSerializer, AmbulanceUpdateSerializer
from ambulance.serializers import CallSerializer

from hospital.models import Hospital, HospitalEquipment
from hospital.serializers import HospitalSerializer, \
    HospitalEquipmentSerializer

logger = logging.getLogger(__name__)

# SubscribeClient

class SubscribeClient(BaseClient):

    # The callback for when the client receives a CONNACK
    # response from the server.
    def on_connect(self, client, userdata, flags, rc):

        # is connected?
        if not super().on_connect(client, userdata, flags, rc):
            return False

        # Subscribing in on_connect() means that if we lose the
        # connection and reconnect then subscriptions will be renewed.
        # client.subscribe('#', 2)

        # ambulance handler
        self.client.message_callback_add('user/+/ambulance/+/data',
                                         self.on_ambulance)

        # hospital handler
        self.client.message_callback_add('user/+/hospital/+/data',
                                         self.on_hospital)

        # hospital equipment handler
        self.client.message_callback_add('user/+/hospital/+/equipment/+/data',
                                         self.on_hospital_equipment)

        # call handler
        # self.client.message_callback_add('ambulance/+/call',
        #                                 self.on_call)

        # client status handler
        self.client.message_callback_add('user/+/client/+/status',
                                         self.on_client_status)

        # subscribe
        self.subscribe('user/+/ambulance/+/data', 2)
        self.subscribe('user/+/hospital/+/data', 2)
        self.subscribe('user/+/hospital/+/equipment/+/data', 2)
        self.subscribe('user/+/client/+/status', 2)

        if self.verbosity > 0:
            self.stdout.write(self.style.SUCCESS(">> Listening to MQTT messages..."))

        return True

    def send_error_message(self, username, topic, payload, error):

        logger.debug("send_error_message: {}, '{}:{}': '{}'".format(username,
                                                                    topic,
                                                                    payload,
                                                                    error))

        if self.verbosity > 0:
            self.stdout.write(self.style.ERROR("*> Error {}, '{}:{}': {}".format(username,
                                                                                topic,
                                                                                payload,
                                                                                error)))

        try:

            message = JSONRenderer().render({
                'topic': topic,
                'payload': payload,
                'error': error
            })
            self.publish('user/{}/error'.format(username), message)

        except Exception as e:

            logger.warning(('mqtt.SubscribeClient: {}, ' +
                            "topic = '{}:{}', " +
                            "error = '{}', " +
                            "exception = {}").format(username,
                                                     topic,
                                                     payload,
                                                     error,
                                                     e))

    def parse_topic(self, msg, json=True):

        if not msg.payload:
            # empty payload
            return

        if self.verbosity > 0:
            self.stdout.write(self.style.SUCCESS(" > Parsing message '{}:{}'".format(msg.topic,
                                                                                msg.payload)))

        # parse topic
        values = msg.topic.split('/')

        # print('values = {}'.format(values))

        try:

            # retrieve user
            # print(User.objects.all())
            user = User.objects.get(username=values[1])

        except ObjectDoesNotExist as e:

            # does not know username
            # cannot send error message to user
            logger.warning(('mqtt.SubscribeClient: {}, ' +
                            "topic = '{}:{}', " +
                            "error = '{}', " +
                            "exception = {}").format(username,
                                                     topic,
                                                     payload,
                                                     error,
                                                     e))
            return

        if json:

            # parse data
            try:

                # Parse data into json dict
                data = JSONParser().parse(BytesIO(msg.payload))

            except Exception as e:

                # send error message to user
                self.send_error_message(user, msg.topic, msg.payload,
                                        "JSON formatted incorrectly")
                return
        else:

            data = msg.payload.decode()

        if len(values) == 5:

            return user, data, values[3]

        elif len(values) == 7:

            return user, data, values[3], values[5]

        else:

            # send error message to user
            # this should never happen because no subscriptions will match
            # topics with different sizes
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Invalid topic")
            return

    # Update ambulance

    def on_ambulance(self, client, userdata, msg):

        logger.debug("on_ambulance: msg = '{}:{}'".format(msg.topic, msg.payload))

        # parse topic
        values = self.parse_topic(msg)
        if not values:
            return

        logger.debug("on_ambulance: values = '{}'".format(values))

        # retrieve parsed values
        user, data, ambulance_id = values

        try:

            logger.debug('ambulance_id = {}'.format(ambulance_id))

            # retrieve ambulance
            ambulance = Ambulance.objects.get(id=ambulance_id)

        except ObjectDoesNotExist:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Ambulance with id '{}' does not exist".format(ambulance_id))
            return

        except Exception as e:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Exception: '{}'".format(e))
            return

        logger.debug('on_ambulance: ambulance = {}'.format(ambulance))

        try:

            is_valid = False
            if isinstance(data, (list, tuple)):

                # update ambulances in bulk
                serializer = AmbulanceUpdateSerializer(data=data,
                                                       many=True,
                                                       partial=True)

                if serializer.is_valid():

                    # save to database
                    serializer.save(ambulance=ambulance, updated_by=user)
                    is_valid = True

            else:

                # update ambulance
                serializer = AmbulanceSerializer(ambulance,
                                                 data=data,
                                                 partial=True)

                if serializer.is_valid():

                    # save to database
                    serializer.save(updated_by=user)
                    is_valid = True

            if not is_valid:

                logger.debug('on_ambulance: INVALID serializer')

                # send error message to user
                self.send_error_message(user, msg.topic, msg.payload,
                                        serializer.errors)

        except Exception as e:

            logger.debug('on_ambulance: serializer EXCEPTION')

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload, e)

        logger.debug('on_ambulance: DONE')

    # Update hospital
    def on_hospital(self, client, userdata, msg):

        logger.debug("on_hospital: msg = '{}:{}'".format(msg.topic, msg.payload))

        # parse topic
        values = self.parse_topic(msg)
        if not values:
            return

        logger.debug("on_hospital: values = '{}'".format(values))

        # retrieve parsed values
        user, data, hospital_id = values

        try:

            # retrieve hospital
            hospital = Hospital.objects.get(id=hospital_id)

        except ObjectDoesNotExist:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Hospital with id '{}' does not exist".format(hospital_id))
            return

        except Exception as e:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Exception: '{}'".format(e))
            return

        logger.debug('on_hospital: hospital = {}'.format(hospital))

        try:

            # update hospital
            serializer = HospitalSerializer(hospital,
                                            data=data,
                                            partial=True)
            if serializer.is_valid():

                logger.debug('on_hospital: valid serializer')

                # save to database
                serializer.save(updated_by=user)

            else:

                logger.debug('on_hospital: INVALID serializer')

                # send error message to user
                self.send_error_message(user, msg.topic, msg.payload,
                                        serializer.errors)

        except Exception as e:

            logger.debug('on_hospital: serializer EXCEPTION')

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload, e)

        logger.debug('on_hospital: DONE')

    # Update hospital equipment
    def on_hospital_equipment(self, client, userdata, msg):

        logger.debug("on_hospital_equipment: msg = '{}:{}'".format(msg.topic, msg.payload))

        # parse topic
        values = self.parse_topic(msg)
        if not values:
            return

        logger.debug("on_hospital_equipment: values = '{}'".format(values))

        # retrieve parsed values
        user, data, hospital_id, equipment_name = values

        try:

            # retrieve hospital equipment
            hospital_equipment = HospitalEquipment.objects.get(hospital=hospital_id,
                                                               equipment__name=equipment_name)

        except ObjectDoesNotExist:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Hospital equipment with hospital id '{}' and name '{}' does not exist".format(
                                        hospital_id, equipment_name))
            return

        except Exception as e:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Exception: '{}'".format(e))
            return

        logger.debug('on_hospital_equipment: equipment = {}'.format(hospital_equipment))

        try:

            # update hospital equipment
            serializer = HospitalEquipmentSerializer(hospital_equipment,
                                                     data=data,
                                                     partial=True)
            if serializer.is_valid():

                logger.debug('on_hospital_equipment: valid serializer')

                # save to database
                serializer.save(updated_by=user)

            else:

                logger.debug('on_hospital_equipment: INVALID serializer')

                # send error message to user
                self.send_error_message(user, msg.topic, msg.payload,
                                        serializer.errors)

        except Exception as e:

            logger.debug('on_hospital_equipment: serializer EXCEPTION')

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload, e)

        logger.debug('on_hospital_equipment: DONE')

    # update client information
    def on_client_status(self, clnt, userdata, msg):

        logger.debug("on_client_status: msg = '{}:{}'".format(msg.topic, msg.payload))

        # parse topic
        values = self.parse_topic(msg, json=False)
        if not values:
            return

        logger.debug("on_client_status: values = '{}'".format(values))

        # retrieve parsed values
        user, status, client_id = values

        # handle status
        if status == 'online':

            # user just logged in, create record
            client = Client(client_id=client_id, user=user, status=status)
            client.save()

            # log operation
            log = ClientLog(client=client, action=status)
            log.save()

        elif status == 'offline' or status == 'disconnected':

            # user client offline or disconnected
            try:

                # retrieve client record first
                client = Client.objects.get(client_id=client_id)

                # is online?
                if client.status != ClientStatus.O.name:

                    # client is not online
                    logger.debug('on_client_status: not online')

                    # send error message to user
                    self.send_error_message(user, msg.topic, msg.payload,
                                            "client '{}' is not online".format(client_id))

                    return

                # update status
                client.status = status
                client.save()

                # log operation
                log = ClientLog(client=client, action=status)
                log.save()

                # clean up mqtt topic
                self.remove_topic('/user/{}/client/{}/status'.format(user.username, client_id))

            except Client.DoesNotExist:

                logger.debug('on_client_status: INVALID client')

                # send error message to user
                self.send_error_message(user, msg.topic, msg.payload,
                                        "client '{}' is not valid".format(client_id))

    # update calls
    def on_call(self, client, userdata, msg):

        logger.debug("on_call: msg = '{}:{}'".format(msg.topic, msg.payload))

        # parse topic
        values = self.parse_topic(msg)
        if not values:
            return

        logger.debug("on_call: values = '{}'".format(values))

        # retrieve parsed values
        user, data, call_id = values

        try:

            logger.debug('call_id = {}'.format(call_id))

            call = Call.objects.get(id=call_id)

        except ObjectDoesNotExist:

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Call with id '{}' does not exist".format(call_id))

            return

        except Exception as e:
            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload,
                                    "Exception: '{}'".format(e))
            return

        logger.debug('on_call: call = {}'.format(call))

        try:

            # update call
            serializer = CallSerializer(call)  # TODO: Find structure

            if serializer.is_valid():
                logger.debug('on_call: vaild serializer')

                # save to database
                serializer.save(updated_by=user)

            else:

                logger.debug('on_call: INVALID serializer')

                # send error message to user
                self.send_error_message(user, msg.topic, msg.payload,
                                        serializer.errors)

        except Exception as e:

            logger.debug('on_call: serializer EXCEPTION')

            # send error message to user
            self.send_error_message(user, msg.topic, msg.payload, e)

        logger.debug('on_call: DONE')
<|endoftext|>


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

from test.test_base import TestBase
from ambry.identity import DatasetNumber
from ambry.orm import Database, Dataset

@unittest.skip
class Test(TestBase):

    def setUp(self):
        from ambry.run import get_runconfig
        from ambry.library import new_library

        rc = get_runconfig()

        self.library = new_library(rc)

        self.dsn = self.library.database.dsn

        self.d_vid = 'd000mptest001'

        # This is needed when calling the MP functions from a non-mp process
        from ambry.bundle.concurrent import alt_init
        alt_init(self.library)

    @classmethod
    def get_rc(self, name='ambry.yaml'):
        import os
        from ambry.run import get_runconfig
        from test import bundlefiles

        def bf_dir(fn):
            return os.path.join(os.path.dirname(bundlefiles.__file__), fn)

        return get_runconfig(bf_dir('ambry.yaml'))

    def new_db_dataset(self, db,  source='source'):
        return db.new_dataset(vid=self.d_vid, source='example.com', dataset='dataset')

    def new_database(self):
        # FIXME: this connection will not be closed properly in a postgres case.
        db = Database(self.dsn)
        db.open()
        return db

    #@unittest.skipIf("not os.getenv('VIRTUAL_ENV')", 'Must have a VIRTUAL_ENV value set to get the config')
    def test_multi_partitions(self):
        """Test creating multip partitions in a multiprocessing run"""
        from ambry.orm.exc import NotFoundError

        l = self.library
        db = l.database

        l.processes = 4

        try:
            ds = db.dataset(self.d_vid)
            db.remove_dataset(ds)
            db.commit()
        except NotFoundError:
            pass

        # print(db.dsn)

        ds = self.new_db_dataset(db, self.d_vid)

        ds.new_table('footable')

        # FIXME! The cast_to_subclass is required before the MP run to create the bundle source File record,
        # otherwise, it is create in the MP run, and you get an Integrity violation.
        l.bundle(ds.vid).cast_to_subclass().commit()

        ds.commit()

        self.assertIsNotNone(db.dataset(self.d_vid))
        self.assertIsNotNone(self.library.bundle(self.d_vid))

        N = 40

        pool = self.library.process_pool()
        args = [ (ds.vid, i) for i in range(1,N+1)]

        names = pool.map(run_mp, args)

        self.assertEqual(N, len(ds.partitions))

        records = list(sorted(o.name for o in ds.partitions))

        self.assertEquals(sorted(names), records)

    def test_unique_object_gen(self):
        """Test creating multip partitions in a multiprocessing run"""
        from ambry.orm.exc import NotFoundError

        l = self.library
        db = l.database

        l.processes = 8

        N = 40

        try:
            ds = db.dataset(self.d_vid)
            db.remove_dataset(ds)
            db.commit()
        except NotFoundError:
            pass

        ds = self.new_db_dataset(db, self.d_vid)

        b = l.bundle(ds.vid)
        # FIXME! The cast_to_subclass is required before the MP run to create the bundle source File record,
        # otherwise, it is create in the MP run, and you get an Integrity violation.
        b = b.cast_to_subclass()
        b.commit()

        self.assertIsNotNone(db.dataset(self.d_vid))
        self.assertIsNotNone(self.library.bundle(self.d_vid))

        pool = self.library.process_pool()

        b.new_table('pre1')
        b.new_table('pre2')
        b.commit()

        self.assertEqual(2, len(ds.tables))

        b.dataset.new_source_table('pre1')
        b.dataset.new_source_table('pre2')
        b.commit()

        self.assertEqual(2, len(ds.source_tables))

        b.dataset.new_partition(table=b.table('pre1'))
        b.dataset.new_partition(table=b.table('pre2'))
        b.commit()

        self.assertEqual(2, len(ds.source_tables))

        for mp_f, ds_attr in ((run_mp_tables, 'tables'),
                              (run_mp_sourcetables, 'source_tables'),
                              (run_mp_partitions, 'partitions')
                              ):

            print('Running: ', ds_attr)

            names = pool.map(mp_f, [(ds.vid, i) for i in range(1, N+1)])

            b.session.refresh(ds)
            self.assertEqual(N+2, len(getattr(ds, ds_attr)))

            records = list(sorted([o.name for o in getattr(ds, ds_attr)]))

            if ds_attr == 'partitions':
                extra_names = ['example.com-dataset-pre1','example.com-dataset-pre2']
            else:
                extra_names = ['pre1','pre2']

            self.assertEquals(sorted(names+extra_names), records)

    @unittest.skipIf("not os.getenv('VIRTUAL_ENV')", 'Must have a VIRTUAL_ENV value set to get the config')
    def test_top_levels(self):

        from ambry.orm.exc import NotFoundError

        l = self.library
        db = l.database

        l.processes = 8

        N = 5

        try:
            ds = db.dataset(self.d_vid)
            db.remove_dataset(ds)
            db.commit()
        except NotFoundError:
            pass

        ds = self.new_db_dataset(db, self.d_vid)

        b = l.bundle(ds.vid)
        # FIXME! The cast_to_subclass is required before the MP run to create the bundle source File record,
        # otherwise, it is create in the MP run, and you get an Integrity violation.
        b = b.cast_to_subclass()
        b.commit()

        pool = self.library.process_pool()

        names = pool.map(run_mp_combined, [(ds.vid, i) for i in range(1, N + 1)])

        from itertools import chain

        names = sorted(chain(*names))

        self.assertEquals(135, len(names))

def run_mp(b,i):
    """Ingest a source, using only arguments that can be pickled, for multiprocessing access"""
    from ambry.identity import PartialPartitionName

    from ambry.identity import PartialPartitionName
    from ambry.orm.partition import Partition

    table = b.table('footable')

    pname = PartialPartitionName(table='footable', segment=i)

    p = b.partitions.new_partition(pname, type=Partition.TYPE.SEGMENT)

    b.commit()

    return p.name

def run_mp_partitions(b,i):
    """Ingest a source, using only arguments that can be pickled, for multiprocessing access"""
    from ambry.identity import PartialPartitionName
    from ambry.orm.partition import Partition

    b, i = args

    table = b.table('obj01') # The run_mp_tables test must come before partitions to make the obj01 table

    p = b.dataset.new_partition(table, segment=i, type=Partition.TYPE.SEGMENT)

    b.commit()

    return p.name

def run_mp_tables(b,i):
    """Ingest a source, using only arguments that can be pickled, for multiprocessing access"""
    from ambry.orm import Table

    b.commit()
    t = b.dataset.new_table(name="obj{:02d}".format(i))
    b.commit()

    if t:
        return t.name
    else:
        return None

def run_mp_sourcetables(b,i):
    """Ingest a source, using only arguments that can be pickled, for multiprocessing access"""

    b.commit()
    t = b.dataset.new_source_table(name="obj{:02d}".format(i))
    b.commit()

    if t:
        return t.name
    else:
        return None

def run_mp_combined(b,i):
    """Ingest a source, using only arguments that can be pickled, for multiprocessing access"""

    ds = b.dataset

    name = "cobj{:02d}".format(i)

    out = []

    for j in range(9):
        name = "cobj{:02d}".format(i*10+j)
        t = ds.new_table(name)
        st = ds.new_source_table(name)
        p = ds.new_partition(t)

        out += [t.name, st.name, p.name]

    b.commit()

    return out

<commit_msg>Multiprocessing tests use proto library. #187.<commit_after># -*- coding: utf-8 -*-
import unittest

from ambry.orm import Database

from test.proto import TestBase

# TODO: This is completely broken because base class changed.

@unittest.skip
class Test(TestBase):

    def setUp(self):
        from ambry.run import get_runconfig
        from ambry.library import new_library

        rc = get_runconfig()

        self.library = new_library(rc)

        self.dsn = self.library.database.dsn

        self.d_vid = 'd000mptest001'

        # This is needed when calling the MP functions from a non-mp process
        from ambry.bundle.concurrent import alt_init
        alt_init(self.library)

    @classmethod
    def get_rc(self, name='ambry.yaml'):
        import os
        from ambry.run import get_runconfig
        from test import bundlefiles

        def bf_dir(fn):
            return os.path.join(os.path.dirname(bundlefiles.__file__), fn)

        return get_runconfig(bf_dir('ambry.yaml'))

    def new_db_dataset(self, db,  source='source'):
        return db.new_dataset(vid=self.d_vid, source='example.com', dataset='dataset')

    def new_database(self):
        # FIXME: this connection will not be closed properly in a postgres case.
        db = Database(self.dsn)
        db.open()
        return db

    #@unittest.skipIf("not os.getenv('VIRTUAL_ENV')", 'Must have a VIRTUAL_ENV value set to get the config')
    def test_multi_partitions(self):
        """Test creating multip partitions in a multiprocessing run"""
        from ambry.orm.exc import NotFoundError

        l = self.library
        db = l.database

        l.processes = 4

        try:
            ds = db.dataset(self.d_vid)
            db.remove_dataset(ds)
            db.commit()
        except NotFoundError:
            pass

        # print(db.dsn)

        ds = self.new_db_dataset(db, self.d_vid)

        ds.new_table('footable')

        # FIXME! The cast_to_subclass is required before the MP run to create the bundle source File record,
        # otherwise, it is create in the MP run, and you get an Integrity violation.
        l.bundle(ds.vid).cast_to_subclass().commit()

        ds.commit()

        self.assertIsNotNone(db.dataset(self.d_vid))
        self.assertIsNotNone(self.library.bundle(self.d_vid))

        N = 40

        pool = self.library.process_pool()
        args = [ (ds.vid, i) for i in range(1,N+1)]

        names = pool.map(run_mp, args)

        self.assertEqual(N, len(ds.partitions))

        records = list(sorted(o.name for o in ds.partitions))

        self.assertEquals(sorted(names), records)

    def test_unique_object_gen(self):
        """Test creating multip partitions in a multiprocessing run"""
        from ambry.orm.exc import NotFoundError

        l = self.library
        db = l.database

        l.processes = 8

        N = 40

        try:
            ds = db.dataset(self.d_vid)
            db.remove_dataset(ds)
            db.commit()
        except NotFoundError:
            pass

        ds = self.new_db_dataset(db, self.d_vid)

        b = l.bundle(ds.vid)
        # FIXME! The cast_to_subclass is required before the MP run to create the bundle source File record,
        # otherwise, it is create in the MP run, and you get an Integrity violation.
        b = b.cast_to_subclass()
        b.commit()

        self.assertIsNotNone(db.dataset(self.d_vid))
        self.assertIsNotNone(self.library.bundle(self.d_vid))

        pool = self.library.process_pool()

        b.new_table('pre1')
        b.new_table('pre2')
        b.commit()

        self.assertEqual(2, len(ds.tables))

        b.dataset.new_source_table('pre1')
        b.dataset.new_source_table('pre2')
        b.commit()

        self.assertEqual(2, len(ds.source_tables))

        b.dataset.new_partition(table=b.table('pre1'))
        b.dataset.new_partition(table=b.table('pre2'))
        b.commit()

        self.assertEqual(2, len(ds.source_tables))

        for mp_f, ds_attr in ((run_mp_tables, 'tables'),
                              (run_mp_sourcetables, 'source_tables'),
                              (run_mp_partitions, 'partitions')
                              ):

            print('Running: ', ds_attr)

            names = pool.map(mp_f, [(ds.vid, i) for i in range(1, N+1)])

            b.session.refresh(ds)
            self.assertEqual(N+2, len(getattr(ds, ds_attr)))

            records = list(sorted([o.name for o in getattr(ds, ds_attr)]))

            if ds_attr == 'partitions':
                extra_names = ['example.com-dataset-pre1','example.com-dataset-pre2']
            else:
                extra_names = ['pre1','pre2']

            self.assertEquals(sorted(names+extra_names), records)

    @unittest.skipIf("not os.getenv('VIRTUAL_ENV')", 'Must have a VIRTUAL_ENV value set to get the config')
    def test_top_levels(self):

        from ambry.orm.exc import NotFoundError

        l = self.library
        db = l.database

        l.processes = 8

        N = 5

        try:
            ds = db.dataset(self.d_vid)
            db.remove_dataset(ds)
            db.commit()
        except NotFoundError:
            pass

        ds = self.new_db_dataset(db, self.d_vid)

        b = l.bundle(ds.vid)
        # FIXME! The cast_to_subclass is required before the MP run to create the bundle source File record,
        # otherwise, it is create in the MP run, and you get an Integrity violation.
        b = b.cast_to_subclass()
        b.commit()

        pool = self.library.process_pool()

        names = pool.map(run_mp_combined, [(ds.vid, i) for i in range(1, N + 1)])

        from itertools import chain

        names = sorted(chain(*names))

        self.assertEquals(135, len(names))

def run_mp(b,i):
    """Ingest a source, using only arguments that can be pickled, for multiprocessing access"""
    from ambry.identity import PartialPartitionName

    from ambry.identity import PartialPartitionName
    from ambry.orm.partition import Partition

    table = b.table('footable')

    pname = PartialPartitionName(table='footable', segment=i)

    p = b.partitions.new_partition(pname, type=Partition.TYPE.SEGMENT)

    b.commit()

    return p.name

def run_mp_partitions(b,i):
    """Ingest a source, using only arguments that can be pickled, for multiprocessing access"""
    from ambry.identity import PartialPartitionName
    from ambry.orm.partition import Partition

    b, i = args

    table = b.table('obj01') # The run_mp_tables test must come before partitions to make the obj01 table

    p = b.dataset.new_partition(table, segment=i, type=Partition.TYPE.SEGMENT)

    b.commit()

    return p.name

def run_mp_tables(b,i):
    """Ingest a source, using only arguments that can be pickled, for multiprocessing access"""
    from ambry.orm import Table

    b.commit()
    t = b.dataset.new_table(name="obj{:02d}".format(i))
    b.commit()

    if t:
        return t.name
    else:
        return None

def run_mp_sourcetables(b,i):
    """Ingest a source, using only arguments that can be pickled, for multiprocessing access"""

    b.commit()
    t = b.dataset.new_source_table(name="obj{:02d}".format(i))
    b.commit()

    if t:
        return t.name
    else:
        return None

def run_mp_combined(b,i):
    """Ingest a source, using only arguments that can be pickled, for multiprocessing access"""

    ds = b.dataset

    name = "cobj{:02d}".format(i)

    out = []

    for j in range(9):
        name = "cobj{:02d}".format(i*10+j)
        t = ds.new_table(name)
        st = ds.new_source_table(name)
        p = ds.new_partition(t)

        out += [t.name, st.name, p.name]

    b.commit()

    return out

<|endoftext|>


<commit_before>"""
Forms for MuckRock
"""

from django import forms
from django.contrib.auth.models import User

import autocomplete_light as autocomplete
from autocomplete_light.contrib.taggit_field import TaggitField, TaggitWidget

from muckrock.agency.models import Agency
from muckrock.jurisdiction.models import Jurisdiction

class MRFilterForm(forms.Form):
    """A generic class to filter a list of items"""
    user = forms.ModelChoiceField(
        required=False,
        queryset=User.objects.all(),
        widget=autocomplete.ChoiceWidget(
            'UserAutocomplete',
            attrs={'placeholder': 'All Users'}))
    agency = forms.ModelChoiceField(
        required=False,
        queryset=Agency.objects.all(),
        widget=autocomplete.ChoiceWidget(
            'AgencyAutocomplete',
            attrs={'placeholder': 'All Agencies'}))
    jurisdiction = forms.ModelChoiceField(
        required=False,
        queryset=Jurisdiction.objects.all(),
        widget=autocomplete.ChoiceWidget(
            'JurisdictionAutocomplete',
            attrs={'placeholder': 'All Jurisdictions'}))
    tags = TaggitField(widget=TaggitWidget(
        'TagAutocomplete',
        attrs={
            'placeholder': 'All Tags',
            'data-autocomplete-minimum-characters': 1}))
<commit_msg>changes placeholder text in filter tags field<commit_after>"""
Forms for MuckRock
"""

from django import forms
from django.contrib.auth.models import User

import autocomplete_light as autocomplete
from autocomplete_light.contrib.taggit_field import TaggitField, TaggitWidget

from muckrock.agency.models import Agency
from muckrock.jurisdiction.models import Jurisdiction

class MRFilterForm(forms.Form):
    """A generic class to filter a list of items"""
    user = forms.ModelChoiceField(
        required=False,
        queryset=User.objects.all(),
        widget=autocomplete.ChoiceWidget(
            'UserAutocomplete',
            attrs={'placeholder': 'All Users'}))
    agency = forms.ModelChoiceField(
        required=False,
        queryset=Agency.objects.all(),
        widget=autocomplete.ChoiceWidget(
            'AgencyAutocomplete',
            attrs={'placeholder': 'All Agencies'}))
    jurisdiction = forms.ModelChoiceField(
        required=False,
        queryset=Jurisdiction.objects.all(),
        widget=autocomplete.ChoiceWidget(
            'JurisdictionAutocomplete',
            attrs={'placeholder': 'All Jurisdictions'}))
    tags = TaggitField(widget=TaggitWidget(
        'TagAutocomplete',
        attrs={
            'placeholder': 'All Tags (comma separated)',
            'data-autocomplete-minimum-characters': 1}))
<|endoftext|>


<commit_before>"""
Tests for site level functionality and helper functions for application tests
"""

from django.core.urlresolvers import reverse
from django.test import TestCase

import logging
import nose.tools

# allow methods that could be functions and too many public methods in tests and **kwarg magic
# pylint: disable=R0201
# pylint: disable=R0904
# pylint: disable=W0142

logging.disable(logging.CRITICAL)

kwargs = {"wsgi.url_scheme": "https"}

# helper functions for view testing
def get_allowed(client, url, templates=None, base='base.html', context=None, redirect=None):
    """Test a get on a url that is allowed with the users current credntials"""
    # pylint: disable=R0913
    response = client.get(url, follow=True, **kwargs)
    nose.tools.eq_(response.status_code, 200)

    if redirect:
        nose.tools.eq_(response.redirect_chain, [('https://testserver:80' + redirect, 302)])

    if templates:
        resp_ts = [t.name for t in response.templates[:3]
                   if t.name != 'compressor/css_file.html']
        nose.tools.eq_(resp_ts, templates + [base])

    if context:
        for key, value in context.iteritems():
            nose.tools.eq_(response.context[key], value)

    return response

def post_allowed(client, url, data, redirect):
    """Test an allowed post with the given data and redirect location"""
    response = client.post(url, data, follow=True, **kwargs)
    nose.tools.eq_(response.status_code, 200)
    nose.tools.eq_(response.redirect_chain, [('https://testserver:80' + redirect, 302)])

    return response

def post_allowed_bad(client, url, templates, data=None):
    """Test an allowed post with bad data"""
    if data is None:
        data = {'bad': 'data'}
    response = client.post(url, data, **kwargs)
    nose.tools.eq_(response.status_code, 200)
    # make sure first 3 match (4th one might be form.html, not important
    nose.tools.eq_([t.name for t in response.templates][:3], templates + ['base.html'])

def get_post_unallowed(client, url):
    """Test an unauthenticated get and post on a url that is allowed
    to be viewed only by authenticated users"""
    redirect = 'https://testserver:80/accounts/login/?next=' + url
    response = client.get(url, **kwargs)
    nose.tools.eq_(response.status_code, 302)
    nose.tools.eq_(response['Location'], redirect)

def get_404(client, url):
    """Test a get on a url that is allowed with the users current credntials"""
    response = client.get(url, **kwargs)
    nose.tools.eq_(response.status_code, 404)

    return response

class TestFunctional(TestCase):
    """Functional tests for top level"""
    fixtures = ['holidays.json', 'jurisdictions.json', 'agency_types.json', 'test_agencies.json',
                'test_users.json', 'test_foiarequests.json', 'test_news.json']

    # tests for base level views
    def test_views(self):
        """Test views"""

        get_allowed(self.client, reverse('index'))
        get_allowed(self.client, '/sitemap.xml')
        get_allowed(self.client, '/search/')
<commit_msg>test fix<commit_after>"""
Tests for site level functionality and helper functions for application tests
"""

from django.core.urlresolvers import reverse
from django.test import TestCase

import logging
import nose.tools

# allow methods that could be functions and too many public methods in tests and **kwarg magic
# pylint: disable=R0201
# pylint: disable=R0904
# pylint: disable=W0142

logging.disable(logging.CRITICAL)

kwargs = {"wsgi.url_scheme": "https"}

# helper functions for view testing
def get_allowed(client, url, templates=None, base='base.html', context=None, redirect=None):
    """Test a get on a url that is allowed with the users current credntials"""
    # pylint: disable=R0913
    response = client.get(url, follow=True, **kwargs)
    nose.tools.eq_(response.status_code, 200)

    if redirect:
        nose.tools.eq_(response.redirect_chain, [('https://testserver:80' + redirect, 302)])

    if templates:
        resp_ts = [t.name for t in response.templates[:len(templates)+1]]
        nose.tools.eq_(resp_ts, templates + [base])

    if context:
        for key, value in context.iteritems():
            nose.tools.eq_(response.context[key], value)

    return response

def post_allowed(client, url, data, redirect):
    """Test an allowed post with the given data and redirect location"""
    response = client.post(url, data, follow=True, **kwargs)
    nose.tools.eq_(response.status_code, 200)
    nose.tools.eq_(response.redirect_chain, [('https://testserver:80' + redirect, 302)])

    return response

def post_allowed_bad(client, url, templates, data=None):
    """Test an allowed post with bad data"""
    if data is None:
        data = {'bad': 'data'}
    response = client.post(url, data, **kwargs)
    nose.tools.eq_(response.status_code, 200)
    # make sure first 3 match (4th one might be form.html, not important
    nose.tools.eq_([t.name for t in response.templates][:3], templates + ['base.html'])

def get_post_unallowed(client, url):
    """Test an unauthenticated get and post on a url that is allowed
    to be viewed only by authenticated users"""
    redirect = 'https://testserver:80/accounts/login/?next=' + url
    response = client.get(url, **kwargs)
    nose.tools.eq_(response.status_code, 302)
    nose.tools.eq_(response['Location'], redirect)

def get_404(client, url):
    """Test a get on a url that is allowed with the users current credntials"""
    response = client.get(url, **kwargs)
    nose.tools.eq_(response.status_code, 404)

    return response

class TestFunctional(TestCase):
    """Functional tests for top level"""
    fixtures = ['holidays.json', 'jurisdictions.json', 'agency_types.json', 'test_agencies.json',
                'test_users.json', 'test_foiarequests.json', 'test_news.json']

    # tests for base level views
    def test_views(self):
        """Test views"""

        get_allowed(self.client, reverse('index'))
        get_allowed(self.client, '/sitemap.xml')
        get_allowed(self.client, '/search/')
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Copyright 2014 RethinkDB, all rights reserved.

"""The `interface.table_acks_durability` test checks that write-acks and durability settings on tables behave as expected."""

from __future__ import print_function

import os, pprint, sys, time

try:
    xrange
except NameError:
    xrange = range

startTime = time.time()

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, 'common')))
import driver, scenario_common, utils, vcoptparse

op = vcoptparse.OptParser()
scenario_common.prepare_option_parser_mode_flags(op)
_, command_prefix, serve_options = scenario_common.parse_mode_flags(op.parse(sys.argv))

r = utils.import_python_driver()
dbName, _ = utils.get_test_db_table()

num_live = num_dead = 3

print("Starting cluster of %d servers (%.2fs)" % (num_live + num_dead, time.time() - startTime))
with driver.Cluster(output_folder='.') as cluster:

    # The "live" processes will remain alive for the entire test. The "dead" processes
    # will be killed after we create some tables.

    live_names = ["l%d" % (i+1) for i in xrange(num_live)]
    dead_names = ["d%d" % (i+1) for i in xrange(num_dead)]

    for name in live_names:
        driver.Process(cluster=cluster, files=name, command_prefix=command_prefix, extra_options=serve_options)

    server = cluster[0]

    dead_procs = []
    for name in dead_names:
        dead_procs.append(driver.Process(cluster=cluster, files=name, command_prefix=command_prefix, extra_options=serve_options))

    cluster.wait_until_ready()
    cluster.check()

    print("Establishing ReQl connections (%.2fs)" % (time.time() - startTime))

    conn = r.connect(host=server.host, port=server.driver_port)

    def mks(nl, nd, primary_replica="l"):
        """Helper function for constructing entries for `table_config.shards`. Returns a
        shard with "nl" live replicas and "nd" dead ones. The value of "primary_replica"
        determines if the primary replica will be a live one or dead one."""
        assert nl <= num_live and nd <= num_dead
        assert (primary_replica == "l" and nl > 0) or (primary_replica == "d" and nd > 0)
        replicas = live_names[:nl] + dead_names[:nd]
        return {"replicas": replicas, "primary_replica": "%s1" % primary_replica}
    def mkr(nl, nd, mode):
        """Helper function for constructing lists for `table_config.write_acks`. Returns
        an ack requirement with "nl" live replicas and "nd" dead ones, and the given mode
        for "single" or "majority"."""
        assert nl <= num_live and nd <= num_dead
        replicas = live_names[:nl] + dead_names[:nd]
        return {"replicas": replicas, "acks": mode}

    # Each test is a tuple with four parts:
    #  1. A list of shards, in the same format as `table_config.shards`.
    #  2. A value in the same format as `table_config.write_acks`.
    #  3. The level of availability we expect for the table after the "dead" servers have
    #     shut down, expressed as a string with the chars `a`, `w`, `r`, and/or `o`,
    #     which correspond to the four fields in the `status` group of `table_status`.
    #  4. The level of availability we expect after the "dead" servers have been
    #     permanently removed and the primary replica (if dead) reassigned to a live server.
    tests = [
        ([mks(3, 0)], "single", "awro", "awro"),
        ([mks(0, 3, "d")], "single", "", ""),
        ([mks(1, 0)], "majority", "awro", "awro"),
        ([mks(0, 1, "d")], "majority", "", ""),
        ([mks(2, 1)], "single", "wro", "awro"),
        ([mks(1, 2)], "single", "wro", "awro"),
        ([mks(2, 1)], "majority", "wro", "awro"),
        ([mks(2, 1, "d")], "majority", "o", "o"),
        ([mks(1, 2)], "majority", "ro", "awro"),
        ([mks(2, 2)], "majority", "ro", "awro"),
        ([mks(2, 2, "d")], "majority", "o", "o"),
        ([mks(2, 2)], [mkr(1, 1, "single")], "wro", "awro"),
        ([mks(2, 2)], [mkr(1, 1, "single"), mkr(2, 0, "majority")], "wro", "awro"),
        ([mks(3, 3)], [mkr(3, 0, "majority"), mkr(0, 3, "majority")], "ro", "ro"),
        ([mks(3, 3)], [mkr(3, 0, "majority"), mkr(0, 1, "single")], "ro", "ro"),
        ([mks(2, 2), mks(2, 2, "d")], [mkr(2, 0, "majority"), mkr(0, 2, "majority")],
            "o", "o"),
        ]

    print("Creating tables for tests (%.2fs)" % (time.time() - startTime))
    r.db_create("test").run(conn)
    for i, (shards, write_acks, readiness_1, readiness_2) in enumerate(tests):
        conf = {"shards": shards, "write_acks": write_acks}
        t = "table%d" % (i+1)
        print(t, conf)
        res = r.db(dbName).table_create(t).run(conn)
        assert res["tables_created"] == 1, res
        res = r.db(dbName).table(t).config().update(conf).run(conn)
        assert res["errors"] == 0, res
        r.db(dbName).table(t).wait().run(conn)
        res = r.db(dbName).table(t).insert([{}]*1000).run(conn)
        assert res["errors"] == 0 and res["inserted"] == 1000, res

    issues = list(r.db("rethinkdb").table("current_issues").run(conn))
    assert not issues, repr(issues)

    print("Killing the designated 'dead' servers (%.2fs)" % (time.time() - startTime))
    for proc in dead_procs:
        proc.check_and_stop()

    print("Checking table statuses (%.2fs)" % (time.time() - startTime))
    def check_table_status(name, expected_readiness):
        tested_readiness = ""
        try:
            res = r.db(dbName).table(name).update({"x": r.row["x"].default(0).add(1)}).run(conn)
        except r.RqlRuntimeError, e:
            # This can happen if we aren't available for reading either
            pass
        else:
            assert res["errors"] == 0 and res["replaced"] == 1000
            tested_readiness += "w"
        try:
            res = r.db(dbName).table(name).count().run(conn)
        except r.RqlRuntimeError, e:
            pass
        else:
            assert res == 1000
            tested_readiness += "r"
        try:
            res = r.db(dbName).table(name).count().run(conn, read_mode="outdated")
        except r.RqlRuntimeError, e:
            pass
        else:
            assert res == 1000
            tested_readiness += "o"
        res = r.db(dbName).table(name).status().run(conn)
        reported_readiness = ""
        if res["status"]["all_replicas_ready"]:
            reported_readiness += "a"
        if res["status"]["ready_for_writes"]:
            reported_readiness += "w"
        if res["status"]["ready_for_reads"]:
            reported_readiness += "r"
        if res["status"]["ready_for_outdated_reads"]:
            reported_readiness += "o"
        print("%s expect=%r test=%r report=%r" % (name, expected_readiness, tested_readiness, reported_readiness))
        assert expected_readiness.replace("a", "") == tested_readiness
        assert expected_readiness == reported_readiness

    for i, (shards, write_acks, readiness_1, readiness_2) in enumerate(tests):
        check_table_status("table%d" % (i+1), readiness_1)

    print("Permanently removing the designated 'dead' servers (%.2fs)" % (time.time() - startTime))
    res = r.db("rethinkdb").table("server_config").filter(lambda s: r.expr(dead_names).contains(s["name"])).delete().run(conn)
    assert res["deleted"] == num_dead, res

    print("Checking table statuses (%.2fs)" % (time.time() - startTime))
    for i, (shards, write_acks, readiness_1, readiness_2) in enumerate(tests):
        check_table_status("table%d" % (i+1), readiness_2)

    print("Checking for issues (%.2fs)" % (time.time() - startTime))
    issues = list(r.db("rethinkdb").table("current_issues").run(conn))
    pprint.pprint(issues)
    issues_by_table = {}
    for issue in issues:
        assert issue["type"] in ["write_acks", "data_lost", "table_needs_primary"]
        issues_by_table[issue["info"]["table"]] = issue["type"]
    for i, (shards, write_acks, readiness_1, readiness_2) in enumerate(tests):
        t = "table%d" % (i+1)
        if readiness_2 == "":
            assert issues_by_table[t] == "data_lost"
        elif readiness_2 == "o":
            assert issues_by_table[t] == "table_needs_primary"
        elif readiness_2 == "ro":
            assert issues_by_table[t] == "write_acks"
        elif readiness_2 == "wro":
            assert False, "This is impossible since we removed the dead servers"
        elif readiness_2 == "awro":
            assert t not in issues_by_table

    print("Running auxiliary tests (%.2fs)" % (time.time() - startTime))
    res = r.db(dbName).table_create("aux").run(conn)
    assert res["tables_created"] == 1, res
    res = r.db(dbName).table("aux").config() \
        .update({"shards": [{"primary_replica": "l1", "replicas": ["l1"]}]}).run(conn)
    assert res["errors"] == 0, res
    r.db(dbName).table("aux").wait().run(conn)
    def test_ok(change):
        print(repr(change))
        res = r.db(dbName).table("aux").config().update(change).run(conn)
        assert res["errors"] == 0 and res["replaced"] == 1, res
        print("OK")
    def test_fail(change):
        print(repr(change))
        res = r.db(dbName).table("aux").config().update(change).run(conn)
        assert res["errors"] == 1 and res["replaced"] == 0, res
        print("Failed (as expected):", repr(res["first_error"]))
    test_ok({"durability": "soft"})
    test_ok({"durability": "hard"})
    test_fail({"durability": "the consistency of toothpaste"})
    test_ok({"write_acks": [{"replicas": ["l1"], "acks": "single"}]})
    test_fail({"write_acks": "this is a string"})
    test_fail({"write_acks": [{"replicas": ["not_a_server"], "acks": "single"}]})
    test_fail({"write_acks": [{"replicas": 1, "acks": "single"}]})
    test_fail({"write_acks": [{"replicas": ["l1"], "acks": 12345}]})
    test_fail({"write_acks": [{"replicas": ["l1"]}]})
    test_fail({"write_acks": [{"acks": "single"}]})
    test_fail({"write_acks": [{"replicas": ["l1"], "acks": "single", "foo": "bar"}]})

    print("Checking that unsatisfiable write acks are rejected (%.2fs)" % (time.time() - startTime))
    test_ok({
        "shards": [
            {"primary_replica": "l1", "replicas": ["l1", "l2", "l3"]},
            {"primary_replica": "l1", "replicas": ["l1"]},
            ],
        "write_acks": "single"
        })
    test_fail({"write_acks": "majority"})

    print("Cleaning up (%.2fs)" % (time.time() - startTime))
print("Done. (%.2fs)" % (time.time() - startTime))

<commit_msg>Removed table_acks_durability sub tests that don't work in 2.1 anymore.<commit_after>#!/usr/bin/env python
# Copyright 2014 RethinkDB, all rights reserved.

"""The `interface.table_acks_durability` test checks that write-acks and durability settings on tables behave as expected."""

from __future__ import print_function

import os, pprint, sys, time

try:
    xrange
except NameError:
    xrange = range

startTime = time.time()

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir, 'common')))
import driver, scenario_common, utils, vcoptparse

op = vcoptparse.OptParser()
scenario_common.prepare_option_parser_mode_flags(op)
_, command_prefix, serve_options = scenario_common.parse_mode_flags(op.parse(sys.argv))

r = utils.import_python_driver()
dbName, _ = utils.get_test_db_table()

num_live = num_dead = 3

print("Starting cluster of %d servers (%.2fs)" % (num_live + num_dead, time.time() - startTime))
with driver.Cluster(output_folder='.') as cluster:

    # The "live" processes will remain alive for the entire test. The "dead" processes
    # will be killed after we create some tables.

    live_names = ["l%d" % (i+1) for i in xrange(num_live)]
    dead_names = ["d%d" % (i+1) for i in xrange(num_dead)]

    for name in live_names:
        driver.Process(cluster=cluster, files=name, command_prefix=command_prefix, extra_options=serve_options)

    server = cluster[0]

    dead_procs = []
    for name in dead_names:
        dead_procs.append(driver.Process(cluster=cluster, files=name, command_prefix=command_prefix, extra_options=serve_options))

    cluster.wait_until_ready()
    cluster.check()

    print("Establishing ReQl connections (%.2fs)" % (time.time() - startTime))

    conn = r.connect(host=server.host, port=server.driver_port)

    def mks(nl, nd, primary_replica="l"):
        """Helper function for constructing entries for `table_config.shards`. Returns a
        shard with "nl" live replicas and "nd" dead ones. The value of "primary_replica"
        determines if the primary replica will be a live one or dead one."""
        assert nl <= num_live and nd <= num_dead
        assert (primary_replica == "l" and nl > 0) or (primary_replica == "d" and nd > 0)
        replicas = live_names[:nl] + dead_names[:nd]
        return {"replicas": replicas, "primary_replica": "%s1" % primary_replica}

    # Each test is a tuple with four parts:
    #  1. A list of shards, in the same format as `table_config.shards`.
    #  2. A value in the same format as `table_config.write_acks`.
    #  3. The level of availability we expect for the table after the "dead" servers have
    #     shut down, expressed as a string with the chars `a`, `w`, `r`, and/or `o`,
    #     which correspond to the four fields in the `status` group of `table_status`.
    #  4. The level of availability we expect after the "dead" servers have been
    #     permanently removed and the primary replica (if dead) reassigned to a live server.
    tests = [
        ([mks(3, 0)], "single", "awro", "awro"),
        ([mks(0, 3, "d")], "single", "", ""),
        ([mks(1, 0)], "majority", "awro", "awro"),
        ([mks(0, 1, "d")], "majority", "", ""),
        ([mks(2, 1)], "single", "wro", "awro"),
        ([mks(1, 2)], "single", "wro", "awro"),
        ([mks(2, 1)], "majority", "wro", "awro"),
        ([mks(2, 1, "d")], "majority", "o", "o"),
        ([mks(1, 2)], "majority", "ro", "awro"),
        ([mks(2, 2)], "majority", "ro", "awro"),
        ([mks(2, 2, "d")], "majority", "o", "o")
        ]

    print("Creating tables for tests (%.2fs)" % (time.time() - startTime))
    r.db_create("test").run(conn)
    for i, (shards, write_acks, readiness_1, readiness_2) in enumerate(tests):
        conf = {"shards": shards, "write_acks": write_acks}
        t = "table%d" % (i+1)
        print(t, conf)
        res = r.db(dbName).table_create(t).run(conn)
        assert res["tables_created"] == 1, res
        res = r.db(dbName).table(t).config().update(conf).run(conn)
        assert res["errors"] == 0, res
        r.db(dbName).table(t).wait().run(conn)
        res = r.db(dbName).table(t).insert([{}]*1000).run(conn)
        assert res["errors"] == 0 and res["inserted"] == 1000, res

    issues = list(r.db("rethinkdb").table("current_issues").run(conn))
    assert not issues, repr(issues)

    print("Killing the designated 'dead' servers (%.2fs)" % (time.time() - startTime))
    for proc in dead_procs:
        proc.check_and_stop()

    print("Checking table statuses (%.2fs)" % (time.time() - startTime))
    def check_table_status(name, expected_readiness):
        tested_readiness = ""
        try:
            res = r.db(dbName).table(name).update({"x": r.row["x"].default(0).add(1)}).run(conn)
        except r.RqlRuntimeError, e:
            # This can happen if we aren't available for reading either
            pass
        else:
            assert res["errors"] == 0 and res["replaced"] == 1000
            tested_readiness += "w"
        try:
            res = r.db(dbName).table(name).count().run(conn)
        except r.RqlRuntimeError, e:
            pass
        else:
            assert res == 1000
            tested_readiness += "r"
        try:
            res = r.db(dbName).table(name).count().run(conn, read_mode="outdated")
        except r.RqlRuntimeError, e:
            pass
        else:
            assert res == 1000
            tested_readiness += "o"
        res = r.db(dbName).table(name).status().run(conn)
        reported_readiness = ""
        if res["status"]["all_replicas_ready"]:
            reported_readiness += "a"
        if res["status"]["ready_for_writes"]:
            reported_readiness += "w"
        if res["status"]["ready_for_reads"]:
            reported_readiness += "r"
        if res["status"]["ready_for_outdated_reads"]:
            reported_readiness += "o"
        print("%s expect=%r test=%r report=%r" % (name, expected_readiness, tested_readiness, reported_readiness))
        assert expected_readiness.replace("a", "") == tested_readiness
        assert expected_readiness == reported_readiness

    for i, (shards, write_acks, readiness_1, readiness_2) in enumerate(tests):
        check_table_status("table%d" % (i+1), readiness_1)

    print("Permanently removing the designated 'dead' servers (%.2fs)" % (time.time() - startTime))
    res = r.db("rethinkdb").table("server_config").filter(lambda s: r.expr(dead_names).contains(s["name"])).delete().run(conn)
    assert res["deleted"] == num_dead, res

    print("Checking table statuses (%.2fs)" % (time.time() - startTime))
    for i, (shards, write_acks, readiness_1, readiness_2) in enumerate(tests):
        check_table_status("table%d" % (i+1), readiness_2)

    print("Checking for issues (%.2fs)" % (time.time() - startTime))
    issues = list(r.db("rethinkdb").table("current_issues").run(conn))
    pprint.pprint(issues)
    issues_by_table = {}
    for issue in issues:
        assert issue["type"] in ["write_acks", "data_lost", "table_needs_primary"]
        issues_by_table[issue["info"]["table"]] = issue["type"]
    for i, (shards, write_acks, readiness_1, readiness_2) in enumerate(tests):
        t = "table%d" % (i+1)
        if readiness_2 == "":
            assert issues_by_table[t] == "data_lost"
        elif readiness_2 == "o":
            assert issues_by_table[t] == "table_needs_primary"
        elif readiness_2 == "ro":
            assert issues_by_table[t] == "write_acks"
        elif readiness_2 == "wro":
            assert False, "This is impossible since we removed the dead servers"
        elif readiness_2 == "awro":
            assert t not in issues_by_table

    print("Running auxiliary tests (%.2fs)" % (time.time() - startTime))
    res = r.db(dbName).table_create("aux").run(conn)
    assert res["tables_created"] == 1, res
    res = r.db(dbName).table("aux").config() \
        .update({"shards": [{"primary_replica": "l1", "replicas": ["l1"]}]}).run(conn)
    assert res["errors"] == 0, res
    r.db(dbName).table("aux").wait().run(conn)
    def test_ok(change):
        print(repr(change))
        res = r.db(dbName).table("aux").config().update(change).run(conn)
        assert res["errors"] == 0 and res["replaced"] == 1, res
        print("OK")
    def test_fail(change):
        print(repr(change))
        res = r.db(dbName).table("aux").config().update(change).run(conn)
        assert res["errors"] == 1 and res["replaced"] == 0, res
        print("Failed (as expected):", repr(res["first_error"]))
    test_ok({"durability": "soft"})
    test_ok({"durability": "hard"})
    test_fail({"durability": "the consistency of toothpaste"})
    test_ok({"write_acks": "single"]})
    test_ok({"write_acks": "majority"]})
    test_fail({"write_acks": [{"replicas": ["l1"], "acks": "single"}]}) # This used to be ok before 2.1
    test_fail({"write_acks": "this is a string"})
    test_fail({"write_acks": [{"replicas": ["not_a_server"], "acks": "single"}]})
    test_fail({"write_acks": [{"replicas": 1, "acks": "single"}]})
    test_fail({"write_acks": [{"replicas": ["l1"], "acks": 12345}]})
    test_fail({"write_acks": [{"replicas": ["l1"]}]})
    test_fail({"write_acks": [{"acks": "single"}]})
    test_fail({"write_acks": [{"replicas": ["l1"], "acks": "single", "foo": "bar"}]})

    print("Checking that unsatisfiable write acks are rejected (%.2fs)" % (time.time() - startTime))
    test_ok({
        "shards": [
            {"primary_replica": "l1", "replicas": ["l1", "l2", "l3"]},
            {"primary_replica": "l1", "replicas": ["l1"]},
            ],
        "write_acks": "single"
        })
    test_fail({"write_acks": "majority"})

    print("Cleaning up (%.2fs)" % (time.time() - startTime))
print("Done. (%.2fs)" % (time.time() - startTime))

<|endoftext|>


<commit_before>"""Plugins support."""

from muffin.utils import LStruct, MuffinException

class PluginException(MuffinException):

    """Implement any exception in plugins."""

class PluginMeta(type):

    """Ensure that each plugin is singleton."""

    _instances = {}

    def __call__(cls, *args, **kwargs):
        """Check for the plugin is initialized already."""
        if not cls.name:
            raise PluginException('Plugin `%s` doesn\'t have a name.' % cls)

        if cls not in cls._instances:
            cls._instances[cls] = super().__call__(*args, **kwargs)
        return cls._instances[cls]

class BasePlugin(metaclass=PluginMeta):

    """Base class for Muffin plugins."""

    # Plugin options with default values
    defaults = {}

    # Plugin dependencies (name: Plugin)
    dependencies = {}

    name = None

    def __init__(self, app=None, **options):
        """Save application and create he plugin's configuration."""
        self.config = self.cfg = LStruct(options)
        self.app = app
        if app:
            app.install(self)

    def setup(self, app):
        """Initialize the plugin.

        Fill the plugin's options from application.
        """
        self.app = app
        for name, ptype in self.dependencies.items():
            if name not in app.ps or not isinstance(app.ps[name], ptype):
                raise PluginException(
                    'Plugin `%s` requires for plugin `%s` to be installed to the application.' % (
                        self.name, ptype))

        for oname, dvalue in self.defaults.items():
            aname = ('%s_%s' % (self.name, oname)).upper()
            app.cfg.setdefault(aname, dvalue)
            self.cfg.setdefault(oname, app.cfg[aname])
<commit_msg>Fix plugins.<commit_after>"""Plugins support."""

from muffin.utils import LStruct, MuffinException

class PluginException(MuffinException):

    """Implement any exception in plugins."""

class PluginMeta(type):

    """Ensure that each plugin is singleton."""

    _instances = {}

    def __call__(cls, *args, **kwargs):
        """Check for the plugin is initialized already."""
        if not cls.name:
            raise PluginException('Plugin `%s` doesn\'t have a name.' % cls)

        if cls not in cls._instances:
            cls._instances[cls] = super().__call__(*args, **kwargs)
        return cls._instances[cls]

class BasePlugin(metaclass=PluginMeta):

    """Base class for Muffin plugins."""

    # Plugin options with default values
    defaults = {}

    # Plugin dependencies (name: Plugin)
    dependencies = {}

    name = None

    def __init__(self, app=None, **options):
        """Save application and create he plugin's configuration."""
        self.config = self.cfg = LStruct(options)
        self.app = app
        if app is not None:
            app.install(self)

    def setup(self, app):
        """Initialize the plugin.

        Fill the plugin's options from application.
        """
        self.app = app
        for name, ptype in self.dependencies.items():
            if name not in app.ps or not isinstance(app.ps[name], ptype):
                raise PluginException(
                    'Plugin `%s` requires for plugin `%s` to be installed to the application.' % (
                        self.name, ptype))

        for oname, dvalue in self.defaults.items():
            aname = ('%s_%s' % (self.name, oname)).upper()
            app.cfg.setdefault(aname, dvalue)
            self.cfg.setdefault(oname, app.cfg[aname])
<|endoftext|>


<commit_before>"""Support session with Muffin framework."""

# Package information
# ===================
import asyncio
import base64
import functools
import time

import ujson as json
from muffin import HTTPFound, Response
from muffin.plugins import BasePlugin
from muffin.utils import create_signature, check_signature, to_coroutine

__version__ = "0.2.0"
__project__ = "muffin-session"
__author__ = "Kirill Klenov <horneds@gmail.com>"
__license__ = "MIT"

SESSION_KEY = 'session'
USER_KEY = 'user'

@asyncio.coroutine
def session_autoload_middleware_factory(app, handler):
    """Create middleware for auto loading sessions."""
    sessions = app.ps.session

    @asyncio.coroutine
    def session_autoload_middleware(request):
        """Load a session to request."""
        yield from sessions.load(request)
        return (yield from handler(request))

    return session_autoload_middleware

class Plugin(BasePlugin):

    """Provide session's engine for Muffin."""

    name = 'session'
    defaults = {
        'auto_load': False,
        'default_user_checker': lambda x: x,
        'login_url': '/login',
        'secret': 'InsecureSecret',
    }

    def setup(self, app):
        """Initialize the plugin."""
        super().setup(app)

        if self.cfg.secret == 'InsecureSecret':
            app.logger.warn(
                'Use insecure secret key. Change SESSION_SECRET option in configuration.')

        self._user_loader = asyncio.coroutine(lambda id_: id_)  # noqa

        app.on_response_prepare.append(self.save)

        if self.cfg.auto_load:
            app.middlewares.append(session_autoload_middleware_factory)

    def user_loader(self, func):
        """Register a function as user loader."""
        self._user_loader = to_coroutine(func)  # noqa
        return self._user_loader

    @asyncio.coroutine
    def load(self, request):
        """Load session from cookies."""
        if SESSION_KEY not in request:
            session = Session(self.cfg.secret)
            session.load(request.cookies)
            self.app.logger.debug('Session loaded: %s', session)
            request[SESSION_KEY] = request.session = session
        return request[SESSION_KEY]

    __call__ = load

    def save(self, request, response):
        """Save session to response cookies."""
        if isinstance(response, Response) and SESSION_KEY in request and not response.started:
            session = request[SESSION_KEY]
            if session.save(response.set_cookie):
                self.app.logger.debug('Session saved: %s', session)

    @asyncio.coroutine
    def load_user(self, request):
        """Load user from request."""
        if USER_KEY not in request:
            session = yield from self.load(request)
            if 'id' not in session:
                return None

            request[USER_KEY] = request.user = yield from self._user_loader(session['id'])

        return request[USER_KEY]

    @asyncio.coroutine
    def check_user(self, request, func=None, location=None, **kwargs):
        """Check for user is logged and pass the given func.

        :param func: user checker function, defaults to default_user_checker
        :param location: where to redirect if user is not logged in.
            May be either string (URL) or function which accepts request as argument
            and returns string URL.
        """
        user = yield from self.load_user(request)
        func = func or self.cfg.default_user_checker
        if not func(user):
            location = location or self.cfg.login_url
            while callable(location):
                location = location(request)
                while asyncio.iscoroutine(location):
                    location = yield from location
            raise HTTPFound(location, **kwargs)
        return user

    def user_pass(self, func=None, location=None, **rkwargs):
        """Decorator ensures that user pass the given func."""
        def wrapper(view):
            view = to_coroutine(view)

            @asyncio.coroutine
            @functools.wraps(view)
            def handler(request, *args, **kwargs):
                yield from self.check_user(request, func, location, **rkwargs)
                return (yield from view(request, *args, **kwargs))
            return handler

        return wrapper

    @asyncio.coroutine
    def login(self, request, id_):
        """Login an user by ID."""
        session = yield from self.load(request)
        session['id'] = id_

    @asyncio.coroutine
    def logout(self, request):
        """Logout an user."""
        session = yield from self.load(request)
        if 'id' in session:
            del session['id']

class Session(dict):

    """Implement session interface."""

    encoding = 'utf-8'

    def __init__(self, secret, key='session.id', **params):
        """Initialize the session."""
        super(Session, self).__init__()
        self.secret = secret.encode(self.encoding)
        self.key = key
        self.params = params
        self.store = {}

    def save(self, set_cookie):
        """Update cookies if the session has been changed."""
        if set(self.store.items()) ^ set(self.items()):
            value = dict(self.items())
            value = json.dumps(value)
            value = self.encrypt(value)
            if not isinstance(value, str):
                value = value.encode(self.encoding)
            set_cookie(self.key, value, **self.params)
            return True
        return False

    def __setitem__(self, name, value):
        """Dump value to JSON."""
        if isinstance(value, (dict, list, tuple, set)):
            value = json.dumps(value)
        super().__setitem__(name, value)

    def load(self, cookies, **kwargs):
        """Load session from cookies."""
        value = cookies.get(self.key, None)
        if value is None:
            return False

        value = self.decrypt(value)
        if not value:
            return False

        data = json.loads(value)
        if not isinstance(data, dict):
            return False

        self.store = data
        self.update(self.store)

    def encrypt(self, value):
        """Encrypt session data."""
        timestamp = str(int(time.time()))
        value = base64.b64encode(value.encode(self.encoding))
        signature = create_signature(self.secret, value + timestamp.encode(),
                                     encoding=self.encoding)
        return "|".join([value.decode(self.encoding), timestamp, signature])

    def decrypt(self, value):
        """Decrypt session data."""
        value, timestamp, signature = value.split("|")
        if check_signature(signature, self.secret, value + timestamp, encoding=self.encoding):
            return base64.b64decode(value).decode(self.encoding)
        return 'null'
<commit_msg>Basic implementation of flash messages<commit_after>"""Support session with Muffin framework."""

# Package information
# ===================
import asyncio
import base64
import functools
import time

import ujson as json
from muffin import HTTPFound, Response
from muffin.plugins import BasePlugin
from muffin.utils import create_signature, check_signature, to_coroutine

__version__ = "0.2.0"
__project__ = "muffin-session"
__author__ = "Kirill Klenov <horneds@gmail.com>"
__license__ = "MIT"

SESSION_KEY = 'session'
USER_KEY = 'user'
FLASHES_KEY = '_flashes'

@asyncio.coroutine
def session_autoload_middleware_factory(app, handler):
    """Create middleware for auto loading sessions."""
    sessions = app.ps.session

    @asyncio.coroutine
    def session_autoload_middleware(request):
        """Load a session to request."""
        yield from sessions.load(request)
        return (yield from handler(request))

    return session_autoload_middleware

class Plugin(BasePlugin):

    """Provide session's engine for Muffin."""

    name = 'session'
    defaults = {
        'auto_load': False,
        'default_user_checker': lambda x: x,
        'login_url': '/login',
        'secret': 'InsecureSecret',
    }

    def setup(self, app):
        """Initialize the plugin."""
        super().setup(app)

        if self.cfg.secret == 'InsecureSecret':
            app.logger.warn(
                'Use insecure secret key. Change SESSION_SECRET option in configuration.')

        self._user_loader = asyncio.coroutine(lambda id_: id_)  # noqa

        app.on_response_prepare.append(self.save)

        if self.cfg.auto_load:
            app.middlewares.append(session_autoload_middleware_factory)

    def user_loader(self, func):
        """Register a function as user loader."""
        self._user_loader = to_coroutine(func)  # noqa
        return self._user_loader

    @asyncio.coroutine
    def load(self, request):
        """Load session from cookies."""
        if SESSION_KEY not in request:
            session = Session(self.cfg.secret)
            session.load(request.cookies)
            self.app.logger.debug('Session loaded: %s', session)
            request[SESSION_KEY] = request.session = session
        return request[SESSION_KEY]

    __call__ = load

    def save(self, request, response):
        """Save session to response cookies."""
        if isinstance(response, Response) and SESSION_KEY in request and not response.started:
            session = request[SESSION_KEY]
            if session.save(response.set_cookie):
                self.app.logger.debug('Session saved: %s', session)

    @asyncio.coroutine
    def load_user(self, request):
        """Load user from request."""
        if USER_KEY not in request:
            session = yield from self.load(request)
            if 'id' not in session:
                return None

            request[USER_KEY] = request.user = yield from self._user_loader(session['id'])

        return request[USER_KEY]

    @asyncio.coroutine
    def check_user(self, request, func=None, location=None, **kwargs):
        """Check for user is logged and pass the given func.

        :param func: user checker function, defaults to default_user_checker
        :param location: where to redirect if user is not logged in.
            May be either string (URL) or function which accepts request as argument
            and returns string URL.
        """
        user = yield from self.load_user(request)
        func = func or self.cfg.default_user_checker
        if not func(user):
            location = location or self.cfg.login_url
            while callable(location):
                location = location(request)
                while asyncio.iscoroutine(location):
                    location = yield from location
            raise HTTPFound(location, **kwargs)
        return user

    def user_pass(self, func=None, location=None, **rkwargs):
        """Decorator ensures that user pass the given func."""
        def wrapper(view):
            view = to_coroutine(view)

            @asyncio.coroutine
            @functools.wraps(view)
            def handler(request, *args, **kwargs):
                yield from self.check_user(request, func, location, **rkwargs)
                return (yield from view(request, *args, **kwargs))
            return handler

        return wrapper

    @asyncio.coroutine
    def login(self, request, id_):
        """Login an user by ID."""
        session = yield from self.load(request)
        session['id'] = id_

    @asyncio.coroutine
    def logout(self, request):
        """Logout an user."""
        session = yield from self.load(request)
        if 'id' in session:
            del session['id']

    @asyncio.coroutine
    def flash(self, request, message, category='message'):
        """
        Add a message to be flashes on next request.
        """
        session = yield from self.load(request)
        flashes = session.get(FLASHES_KEY, [])
        if isinstance(flashes, str):
            flashes = json.loads(flashes)
        flashes.append((category, message))
        session[FLASHES_KEY] = flashes

    def get_flashed_messages(request, with_categories=False, category_fileter=[]):
        """
        Retrieve flashed messages.
        Important: this is not a coroutine,
        and as such it requires you to call `app.ps.session.load` beforehand!
        For example, if you call `load_user` then it is okay.
        """
        if SESSION_KEY not in request:
            raise ValueError('Please call '
                             '`await app.ps.session.load(request)` beforehand!')
        session = request[SESSION_KEY]
        flashes = session.pop(FLASHES_KEY, [])
        if isinstance(flashes, str):
            flashes = json.loads(flashes)
        if category_filter:
            flashes = list(filter(lambda f: f[0] in category_filter, flashes))
        if not with_categories:
            return [f[1] for f in flashes]
        return flashes

class Session(dict):

    """Implement session interface."""

    encoding = 'utf-8'

    def __init__(self, secret, key='session.id', **params):
        """Initialize the session."""
        super(Session, self).__init__()
        self.secret = secret.encode(self.encoding)
        self.key = key
        self.params = params
        self.store = {}

    def save(self, set_cookie):
        """Update cookies if the session has been changed."""
        if set(self.store.items()) ^ set(self.items()):
            value = dict(self.items())
            value = json.dumps(value)
            value = self.encrypt(value)
            if not isinstance(value, str):
                value = value.encode(self.encoding)
            set_cookie(self.key, value, **self.params)
            return True
        return False

    def __setitem__(self, name, value):
        """Dump value to JSON."""
        if isinstance(value, (dict, list, tuple, set)):
            value = json.dumps(value)
        super().__setitem__(name, value)

    def load(self, cookies, **kwargs):
        """Load session from cookies."""
        value = cookies.get(self.key, None)
        if value is None:
            return False

        value = self.decrypt(value)
        if not value:
            return False

        data = json.loads(value)
        if not isinstance(data, dict):
            return False

        self.store = data
        self.update(self.store)

    def encrypt(self, value):
        """Encrypt session data."""
        timestamp = str(int(time.time()))
        value = base64.b64encode(value.encode(self.encoding))
        signature = create_signature(self.secret, value + timestamp.encode(),
                                     encoding=self.encoding)
        return "|".join([value.decode(self.encoding), timestamp, signature])

    def decrypt(self, value):
        """Decrypt session data."""
        value, timestamp, signature = value.split("|")
        if check_signature(signature, self.secret, value + timestamp, encoding=self.encoding):
            return base64.b64decode(value).decode(self.encoding)
        return 'null'
<|endoftext|>