<commit_before># Copyright Hybrid Logic Ltd.  See LICENSE file for details.

"""
Tests for the control service REST API.
"""

import socket

from uuid import uuid4

from twisted.trial.unittest import TestCase

from twisted.web.http import BAD_REQUEST

from treq import get, json_content

from ..testtools import loop_until, random_name
from .testtools import (
    MONGO_IMAGE, require_mongo, get_mongo_client,
)
from ..node.agents.test.test_blockdevice import REALISTIC_BLOCKDEVICE_SIZE
from .testtools import get_test_cluster, require_cluster

def verify_socket(host, port):
    """
    Wait until the destionation can be connected to.

    :param bytes host: Host to connect to.
    :param int port: Port to connect to.

    :return Deferred: Firing when connection is possible.
    """
    def can_connect():
        s = socket.socket()
        conn = s.connect_ex((host, port))
        return False if conn else True

    dl = loop_until(can_connect)
    return dl

class ContainerAPITests(TestCase):
    """
    Tests for the container API.
    """
    def _create_container(self):
        """
        Create a container listening on port 8080.

        :return: ``Deferred`` firing with a tuple of ``Cluster`` instance
        and container dictionary once the container is up and running.
        """
        data = {
            u"name": random_name(),
            u"host": None,
            u"image": "clusterhq/flask:latest",
            u"ports": [{u"internal": 80, u"external": 8080}],
            u'restart_policy': {u'name': u'never'}
        }
        waiting_for_cluster = get_test_cluster(node_count=1)

        def create_container(cluster, data):
            data[u"host"] = cluster.nodes[0].address
            return cluster.create_container(data)

        d = waiting_for_cluster.addCallback(create_container, data)

        def check_result(result):
            cluster, response = result
            self.addCleanup(cluster.remove_container, data[u"name"])

            self.assertEqual(response, data)
            dl = verify_socket(data[u"host"], 8080)
            dl.addCallback(lambda _: (cluster, response))
            return dl

        d.addCallback(check_result)
        return d

    @require_cluster(1)
    def test_create_container_with_ports(self, cluster):
        """
        Create a container including port mappings on a single-node cluster.
        """
        return self._create_container()

    @require_cluster(1)
    def test_create_container_with_environment(self, cluster):
        """
        Create a container including environment variables on a single-node
        cluster.
        """
        data = {
            u"name": random_name(),
            u"host": None,
            u"image": "clusterhq/flaskenv:latest",
            u"ports": [{u"internal": 8080, u"external": 8081}],
            u"environment": {u"ACCEPTANCE_ENV_LABEL": 'acceptance test ok'},
            u'restart_policy': {u'name': u'never'},
        }
        waiting_for_cluster = get_test_cluster(node_count=1)

        def create_container(cluster, data):
            data[u"host"] = cluster.nodes[0].address
            return cluster.create_container(data)

        d = waiting_for_cluster.addCallback(create_container, data)

        def check_result((cluster, response)):
            self.addCleanup(cluster.remove_container, data[u"name"])
            self.assertEqual(response, data)

        def query_environment(host, port):
            """
            The running container, clusterhq/flaskenv, is a simple Flask app
            that returns a JSON dump of the container's environment, so we
            make an HTTP request and parse the response.
            """
            req = get(
                "http://{host}:{port}".format(host=host, port=port),
                persistent=False
            ).addCallback(json_content)
            return req

        d.addCallback(check_result)
        d.addCallback(lambda _: verify_socket(data[u"host"], 8081))
        d.addCallback(lambda _: query_environment(data[u"host"], 8081))
        d.addCallback(
            lambda response:
                self.assertDictContainsSubset(data[u"environment"], response)
        )
        return d

    @require_mongo
    @require_cluster(2)
    def test_move_container_with_dataset(self, cluster):
        """
        Create a mongodb container with an attached dataset, issue API call
        to move the container. Wait until we can connect to the running
        container on the new host and verify the data has moved with it.
        """
        creating_dataset = create_dataset(self, nodes=2)

        def created_dataset(result):
            cluster, dataset = result
            mongodb = {
                u"name": random_name(),
                u"host": cluster.nodes[0].address,
                u"image": MONGO_IMAGE,
                u"ports": [{u"internal": 27017, u"external": 27017}],
                u'restart_policy': {u'name': u'never'},
                u"volumes": [{u"dataset_id": dataset[u"dataset_id"],
                              u"mountpoint": u"/data/db"}],
            }
            created = cluster.create_container(mongodb)
            created.addCallback(lambda _: self.addCleanup(
                cluster.remove_container, mongodb[u"name"]))
            created.addCallback(
                lambda _: get_mongo_client(cluster.nodes[0].address))

            def got_mongo_client(client):
                database = client.example
                database.posts.insert({u"the data": u"it moves"})
                return database.posts.find_one()
            created.addCallback(got_mongo_client)

            def inserted(record):
                moved = cluster.move_container(
                    mongodb[u"name"], cluster.nodes[1].address
                )

                def destroy_and_recreate(_, record):
                    """
                    After moving our container via the API, we then remove the
                    container on the new host and recreate it, pointing to the
                    same dataset, but with the new container instance exposing
                    a different external port. This technique ensures that the
                    test does not pass by mere accident without the container
                    having moved; by recreating the container on its new host
                    after moving, we can be sure that if we can still connect
                    and read the data, the dataset was successfully moved along
                    with the container.
                    """
                    removed = cluster.remove_container(mongodb[u"name"])
                    mongodb2 = mongodb.copy()
                    mongodb2[u"ports"] = [
                        {u"internal": 27017, u"external": 27018}
                    ]
                    mongodb2[u"host"] = cluster.nodes[1].address
                    removed.addCallback(
                        lambda _: cluster.create_container(mongodb2))
                    removed.addCallback(lambda _: record)
                    return removed
                moved.addCallback(destroy_and_recreate, record)
                return moved
            created.addCallback(inserted)

            def moved(record):
                d = get_mongo_client(cluster.nodes[1].address, 27018)
                d.addCallback(lambda client: client.example.posts.find_one())
                d.addCallback(self.assertEqual, record)
                return d

            created.addCallback(moved)
            return created
        creating_dataset.addCallback(created_dataset)
        return creating_dataset

    @require_mongo
    @require_cluster(1)
    def test_create_container_with_dataset(self, cluster):
        """
        Create a mongodb container with an attached dataset, insert some data,
        shut it down, create a new container with same dataset, make sure
        the data is still there.
        """
        creating_dataset = create_dataset(self)

        def created_dataset(result):
            cluster, dataset = result
            mongodb = {
                u"name": random_name(),
                u"host": cluster.nodes[0].address,
                u"image": MONGO_IMAGE,
                u"ports": [{u"internal": 27017, u"external": 27017}],
                u'restart_policy': {u'name': u'never'},
                u"volumes": [{u"dataset_id": dataset[u"dataset_id"],
                              u"mountpoint": u"/data/db"}],
            }
            created = cluster.create_container(mongodb)
            created.addCallback(lambda _: self.addCleanup(
                cluster.remove_container, mongodb[u"name"]))
            created.addCallback(
                lambda _: get_mongo_client(cluster.nodes[0].address))

            def got_mongo_client(client):
                database = client.example
                database.posts.insert({u"the data": u"it moves"})
                return database.posts.find_one()
            created.addCallback(got_mongo_client)

            def inserted(record):
                removed = cluster.remove_container(mongodb[u"name"])
                mongodb2 = mongodb.copy()
                mongodb2[u"ports"] = [{u"internal": 27017, u"external": 27018}]
                removed.addCallback(
                    lambda _: cluster.create_container(mongodb2))
                removed.addCallback(lambda _: record)
                return removed
            created.addCallback(inserted)

            def restarted(record):
                d = get_mongo_client(cluster.nodes[0].address, 27018)
                d.addCallback(lambda client: client.example.posts.find_one())
                d.addCallback(self.assertEqual, record)
                return d
            created.addCallback(restarted)
            return created
        creating_dataset.addCallback(created_dataset)
        return creating_dataset

    @require_cluster(1)
    def test_current(self, cluster):
        """
        The current container endpoint includes a currently running container.
        """
        creating = self._create_container()

        def created(result):
            cluster, data = result
            data[u"running"] = True

            def in_current():
                current = cluster.current_containers()
                current.addCallback(lambda result: data in result[1])
                return current
            return loop_until(in_current)
        creating.addCallback(created)
        return creating

def create_dataset(test_case, nodes=1,
                   maximum_size=REALISTIC_BLOCKDEVICE_SIZE):
    """
    Create a dataset on a cluster.

    :param TestCase test_case: The test the API is running on.
    :param int nodes: The number of nodes to create. Defaults to 1.
    :param int maximum_size: The size of the dataset to create on the test
        cluster.
    :return: ``Deferred`` firing with a tuple of (``Cluster``
        instance, dataset dictionary) once the dataset is present in
        actual cluster state.
    """
    # Create a cluster
    waiting_for_cluster = get_test_cluster(node_count=nodes)

    # Configure a dataset on node1
    def configure_dataset(cluster):
        """
        Send a dataset creation request on node1.
        """
        requested_dataset = {
            u"primary": cluster.nodes[0].address,
            u"dataset_id": unicode(uuid4()),
            u"maximum_size": maximum_size,
            u"metadata": {u"name": u"my_volume"},
        }

        d = cluster.create_dataset(requested_dataset)

        def got_result(result):
            test_case.addCleanup(
                cluster.delete_dataset, requested_dataset[u"dataset_id"])
            return result
        d.addCallback(got_result)
        return d

    configuring_dataset = waiting_for_cluster.addCallback(
        configure_dataset
    )

    # Wait for the dataset to be created
    waiting_for_create = configuring_dataset.addCallback(
        lambda (cluster, dataset): cluster.wait_for_dataset(dataset)
    )

    return waiting_for_create

class DatasetAPITests(TestCase):
    """
    Tests for the dataset API.
    """
    def test_dataset_creation(self):
        """
        A dataset can be created on a specific node.
        """
        return create_dataset(self)

    def test_dataset_move(self):
        """
        A dataset can be moved from one node to another.
        """
        # Create a 2 node cluster
        waiting_for_cluster = get_test_cluster(node_count=2)

        # Configure a dataset on node1
        def configure_dataset(cluster):
            """
            Send a dataset creation request on node1.
            """
            requested_dataset = {
                u"primary": cluster.nodes[0].address,
                u"dataset_id": unicode(uuid4()),
                u"metadata": {u"name": u"my_volume"}
            }

            return cluster.create_dataset(requested_dataset)
        configuring_dataset = waiting_for_cluster.addCallback(
            configure_dataset
        )

        # Wait for the dataset to be created
        waiting_for_create = configuring_dataset.addCallback(
            lambda (cluster, dataset): cluster.wait_for_dataset(dataset)
        )

        # Once created, request to move the dataset to node2
        def move_dataset((cluster, dataset)):
            moved_dataset = {
                u'primary': cluster.nodes[1].address
            }
            return cluster.update_dataset(dataset['dataset_id'], moved_dataset)
        dataset_moving = waiting_for_create.addCallback(move_dataset)

        # Wait for the dataset to be moved
        waiting_for_move = dataset_moving.addCallback(
            lambda (cluster, dataset): cluster.wait_for_dataset(dataset)
        )

        return waiting_for_move

    def test_dataset_deletion(self):
        """
        A dataset can be deleted, resulting in its removal from the node.
        """
        created = create_dataset(self)

        def delete_dataset(result):
            cluster, dataset = result
            deleted = cluster.delete_dataset(dataset["dataset_id"])

            def not_exists():
                request = cluster.datasets_state()
                request.addCallback(
                    lambda actual_datasets: dataset["dataset_id"] not in
                    (d["dataset_id"] for d in actual_datasets))
                return request
            deleted.addCallback(lambda _: loop_until(not_exists))
            return deleted
        created.addCallback(delete_dataset)
        return created

    def test_dataset_grow(self):
        """
        The size of a dataset can be increased.
        """
        creating = create_dataset(
            test_case=self, maximum_size=REALISTIC_BLOCKDEVICE_SIZE
        )
        new_size = REALISTIC_BLOCKDEVICE_SIZE * 2

        def resize_dataset(result):
            cluster, dataset = result
            return cluster.update_dataset(
                dataset["dataset_id"], {u'maximum_size': new_size}
            )

        resizing = creating.addCallback(resize_dataset)

        def check_dataset_size(result):
            cluster, dataset = result
            self.assertEqual(new_size, dataset['maximum_size'])
            return cluster.wait_for_dataset(dataset)

        checking = resizing.addCallback(check_dataset_size)

        return checking

    def test_dataset_shrink(self):
        """
        The size of a dataset can be decreased.
        """
        creating = create_dataset(
            test_case=self, maximum_size=REALISTIC_BLOCKDEVICE_SIZE * 2
        )
        new_size = REALISTIC_BLOCKDEVICE_SIZE

        def resize_dataset(result):
            cluster, dataset = result
            return cluster.update_dataset(
                dataset["dataset_id"], {u'maximum_size': new_size}
            )

        resizing = creating.addCallback(resize_dataset)

        def check_dataset_size(result):
            cluster, dataset = result
            self.assertEqual(new_size, dataset['maximum_size'])
            return cluster.wait_for_dataset(dataset)

        checking = resizing.addCallback(check_dataset_size)

        return checking

    def test_dataset_shrink_not_valid(self):
        """
        If the requested maximum_size is smaller than the allowed minimum the
        response is ``BAD_REQUEST``.
        """
        creating = create_dataset(
            test_case=self, maximum_size=REALISTIC_BLOCKDEVICE_SIZE
        )
        new_size = 67108864 - 1

        def resize_dataset(result):
            cluster, dataset = result
            # Reconfigure that dataset to be an invalid size.
            resizing = cluster.update_dataset(
                dataset_id=dataset["dataset_id"],
                dataset_properties={u'maximum_size': new_size}
            )
            # Check for expected exception and response code.
            return self.assertFailure(
                resizing, ValueError
            ).addCallback(
                lambda exception: self.assertEqual(
                    BAD_REQUEST, exception.args[1]
                )
            )

        return creating.addCallback(resize_dataset)

    def test_dataset_remove_size_limit(self):
        """
        A dataset with a size limit can have that limit removed.
        """
        creating = create_dataset(
            test_case=self, maximum_size=REALISTIC_BLOCKDEVICE_SIZE
        )
        new_size = None

        def resize_dataset(result):
            cluster, dataset = result
            return cluster.update_dataset(
                dataset["dataset_id"], {u'maximum_size': new_size}
            )

        resizing = creating.addCallback(resize_dataset)

        def check_dataset_size(result):
            cluster, dataset = result
            # If there is no maximum_size, the configuration response will not
            # contain that key
            self.assertNotIn(u'maximum_size', dataset.keys())
            return cluster.wait_for_dataset(dataset)

        checking = resizing.addCallback(check_dataset_size)

        return checking
<commit_msg>How are we going to run the acceptance tests against this openstack api?<commit_after># Copyright Hybrid Logic Ltd.  See LICENSE file for details.

"""
Tests for the control service REST API.
"""

import socket

from uuid import uuid4

from twisted.trial.unittest import TestCase

from twisted.web.http import BAD_REQUEST

from treq import get, json_content

from ..testtools import loop_until, random_name
from .testtools import (
    MONGO_IMAGE, require_mongo, get_mongo_client,
)
from ..node.agents.test.test_blockdevice import REALISTIC_BLOCKDEVICE_SIZE
from .testtools import get_test_cluster, require_cluster

def verify_socket(host, port):
    """
    Wait until the destionation can be connected to.

    :param bytes host: Host to connect to.
    :param int port: Port to connect to.

    :return Deferred: Firing when connection is possible.
    """
    def can_connect():
        s = socket.socket()
        conn = s.connect_ex((host, port))
        return False if conn else True

    dl = loop_until(can_connect)
    return dl

class ContainerAPITests(TestCase):
    """
    Tests for the container API.
    """
    def _create_container(self):
        """
        Create a container listening on port 8080.

        :return: ``Deferred`` firing with a tuple of ``Cluster`` instance
        and container dictionary once the container is up and running.
        """
        data = {
            u"name": random_name(),
            u"host": None,
            u"image": "clusterhq/flask:latest",
            u"ports": [{u"internal": 80, u"external": 8080}],
            u'restart_policy': {u'name': u'never'}
        }
        waiting_for_cluster = get_test_cluster(node_count=1)

        def create_container(cluster, data):
            data[u"host"] = cluster.nodes[0].address
            return cluster.create_container(data)

        d = waiting_for_cluster.addCallback(create_container, data)

        def check_result(result):
            cluster, response = result
            self.addCleanup(cluster.remove_container, data[u"name"])

            self.assertEqual(response, data)
            dl = verify_socket(data[u"host"], 8080)
            dl.addCallback(lambda _: (cluster, response))
            return dl

        d.addCallback(check_result)
        return d

    @require_cluster(1)
    def test_create_container_with_ports(self, cluster):
        """
        Create a container including port mappings on a single-node cluster.
        """
        return self._create_container()

    @require_cluster(1)
    def test_create_container_with_environment(self, cluster):
        """
        Create a container including environment variables on a single-node
        cluster.
        """
        data = {
            u"name": random_name(),
            u"host": None,
            u"image": "clusterhq/flaskenv:latest",
            u"ports": [{u"internal": 8080, u"external": 8081}],
            u"environment": {u"ACCEPTANCE_ENV_LABEL": 'acceptance test ok'},
            u'restart_policy': {u'name': u'never'},
        }
        waiting_for_cluster = get_test_cluster(node_count=1)

        def create_container(cluster, data):
            data[u"host"] = cluster.nodes[0].address
            return cluster.create_container(data)

        d = waiting_for_cluster.addCallback(create_container, data)

        def check_result((cluster, response)):
            self.addCleanup(cluster.remove_container, data[u"name"])
            self.assertEqual(response, data)

        def query_environment(host, port):
            """
            The running container, clusterhq/flaskenv, is a simple Flask app
            that returns a JSON dump of the container's environment, so we
            make an HTTP request and parse the response.
            """
            req = get(
                "http://{host}:{port}".format(host=host, port=port),
                persistent=False
            ).addCallback(json_content)
            return req

        d.addCallback(check_result)
        d.addCallback(lambda _: verify_socket(data[u"host"], 8081))
        d.addCallback(lambda _: query_environment(data[u"host"], 8081))
        d.addCallback(
            lambda response:
                self.assertDictContainsSubset(data[u"environment"], response)
        )
        return d

    @require_mongo
    @require_cluster(2)
    def test_move_container_with_dataset(self, cluster):
        """
        Create a mongodb container with an attached dataset, issue API call
        to move the container. Wait until we can connect to the running
        container on the new host and verify the data has moved with it.
        """
        creating_dataset = create_dataset(self, nodes=2)

        def created_dataset(result):
            cluster, dataset = result
            mongodb = {
                u"name": random_name(),
                u"host": cluster.nodes[0].address,
                u"image": MONGO_IMAGE,
                u"ports": [{u"internal": 27017, u"external": 27017}],
                u'restart_policy': {u'name': u'never'},
                u"volumes": [{u"dataset_id": dataset[u"dataset_id"],
                              u"mountpoint": u"/data/db"}],
            }
            created = cluster.create_container(mongodb)
            created.addCallback(lambda _: self.addCleanup(
                cluster.remove_container, mongodb[u"name"]))
            created.addCallback(
                lambda _: get_mongo_client(cluster.nodes[0].address))

            def got_mongo_client(client):
                database = client.example
                database.posts.insert({u"the data": u"it moves"})
                return database.posts.find_one()
            created.addCallback(got_mongo_client)

            def inserted(record):
                moved = cluster.move_container(
                    mongodb[u"name"], cluster.nodes[1].address
                )

                def destroy_and_recreate(_, record):
                    """
                    After moving our container via the API, we then remove the
                    container on the new host and recreate it, pointing to the
                    same dataset, but with the new container instance exposing
                    a different external port. This technique ensures that the
                    test does not pass by mere accident without the container
                    having moved; by recreating the container on its new host
                    after moving, we can be sure that if we can still connect
                    and read the data, the dataset was successfully moved along
                    with the container.
                    """
                    removed = cluster.remove_container(mongodb[u"name"])
                    mongodb2 = mongodb.copy()
                    mongodb2[u"ports"] = [
                        {u"internal": 27017, u"external": 27018}
                    ]
                    mongodb2[u"host"] = cluster.nodes[1].address
                    removed.addCallback(
                        lambda _: cluster.create_container(mongodb2))
                    removed.addCallback(lambda _: record)
                    return removed
                moved.addCallback(destroy_and_recreate, record)
                return moved
            created.addCallback(inserted)

            def moved(record):
                d = get_mongo_client(cluster.nodes[1].address, 27018)
                d.addCallback(lambda client: client.example.posts.find_one())
                d.addCallback(self.assertEqual, record)
                return d

            created.addCallback(moved)
            return created
        creating_dataset.addCallback(created_dataset)
        return creating_dataset

    @require_mongo
    @require_cluster(1)
    def test_create_container_with_dataset(self, cluster):
        """
        Create a mongodb container with an attached dataset, insert some data,
        shut it down, create a new container with same dataset, make sure
        the data is still there.
        """
        creating_dataset = create_dataset(self)

        def created_dataset(result):
            cluster, dataset = result
            mongodb = {
                u"name": random_name(),
                u"host": cluster.nodes[0].address,
                u"image": MONGO_IMAGE,
                u"ports": [{u"internal": 27017, u"external": 27017}],
                u'restart_policy': {u'name': u'never'},
                u"volumes": [{u"dataset_id": dataset[u"dataset_id"],
                              u"mountpoint": u"/data/db"}],
            }
            created = cluster.create_container(mongodb)
            created.addCallback(lambda _: self.addCleanup(
                cluster.remove_container, mongodb[u"name"]))
            created.addCallback(
                lambda _: get_mongo_client(cluster.nodes[0].address))

            def got_mongo_client(client):
                database = client.example
                database.posts.insert({u"the data": u"it moves"})
                return database.posts.find_one()
            created.addCallback(got_mongo_client)

            def inserted(record):
                removed = cluster.remove_container(mongodb[u"name"])
                mongodb2 = mongodb.copy()
                mongodb2[u"ports"] = [{u"internal": 27017, u"external": 27018}]
                removed.addCallback(
                    lambda _: cluster.create_container(mongodb2))
                removed.addCallback(lambda _: record)
                return removed
            created.addCallback(inserted)

            def restarted(record):
                d = get_mongo_client(cluster.nodes[0].address, 27018)
                d.addCallback(lambda client: client.example.posts.find_one())
                d.addCallback(self.assertEqual, record)
                return d
            created.addCallback(restarted)
            return created
        creating_dataset.addCallback(created_dataset)
        return creating_dataset

    @require_cluster(1)
    def test_current(self, cluster):
        """
        The current container endpoint includes a currently running container.
        """
        creating = self._create_container()

        def created(result):
            cluster, data = result
            data[u"running"] = True

            def in_current():
                current = cluster.current_containers()
                current.addCallback(lambda result: data in result[1])
                return current
            return loop_until(in_current)
        creating.addCallback(created)
        return creating

def create_dataset(test_case, nodes=1,
                   maximum_size=REALISTIC_BLOCKDEVICE_SIZE):
    """
    Create a dataset on a cluster.

    :param TestCase test_case: The test the API is running on.
    :param int nodes: The number of nodes to create. Defaults to 1.
    :param int maximum_size: The size of the dataset to create on the test
        cluster.
    :return: ``Deferred`` firing with a tuple of (``Cluster``
        instance, dataset dictionary) once the dataset is present in
        actual cluster state.
    """
    # Create a cluster
    waiting_for_cluster = get_test_cluster(node_count=nodes)

    # Configure a dataset on node1
    def configure_dataset(cluster):
        """
        Send a dataset creation request on node1.
        """
        requested_dataset = {
            u"primary": cluster.nodes[0].address,
            u"dataset_id": unicode(uuid4()),
            u"maximum_size": maximum_size,
            u"metadata": {u"name": u"my_volume"},
        }

        d = cluster.create_dataset(requested_dataset)

        def got_result(result):
            test_case.addCleanup(
                cluster.delete_dataset, requested_dataset[u"dataset_id"])
            return result
        d.addCallback(got_result)
        return d

    configuring_dataset = waiting_for_cluster.addCallback(
        configure_dataset
    )

    # Wait for the dataset to be created
    waiting_for_create = configuring_dataset.addCallback(
        lambda (cluster, dataset): cluster.wait_for_dataset(dataset)
    )

    return waiting_for_create

class DatasetAPITests(TestCase):
    """
    Tests for the dataset API.

    XXX: We'll need a way to run these acceptance tests on a cluster running an OpenStack based dataset convergence agent.
    Maybe something temporary in this branch until we implement https://clusterhq.atlassian.net/browse/FLOC-1420 "Choose backend at install time"
    """
    def test_dataset_creation(self):
        """
        A dataset can be created on a specific node.
        """
        return create_dataset(self)

    def test_dataset_move(self):
        """
        A dataset can be moved from one node to another.
        """
        # Create a 2 node cluster
        waiting_for_cluster = get_test_cluster(node_count=2)

        # Configure a dataset on node1
        def configure_dataset(cluster):
            """
            Send a dataset creation request on node1.
            """
            requested_dataset = {
                u"primary": cluster.nodes[0].address,
                u"dataset_id": unicode(uuid4()),
                u"metadata": {u"name": u"my_volume"}
            }

            return cluster.create_dataset(requested_dataset)
        configuring_dataset = waiting_for_cluster.addCallback(
            configure_dataset
        )

        # Wait for the dataset to be created
        waiting_for_create = configuring_dataset.addCallback(
            lambda (cluster, dataset): cluster.wait_for_dataset(dataset)
        )

        # Once created, request to move the dataset to node2
        def move_dataset((cluster, dataset)):
            moved_dataset = {
                u'primary': cluster.nodes[1].address
            }
            return cluster.update_dataset(dataset['dataset_id'], moved_dataset)
        dataset_moving = waiting_for_create.addCallback(move_dataset)

        # Wait for the dataset to be moved
        waiting_for_move = dataset_moving.addCallback(
            lambda (cluster, dataset): cluster.wait_for_dataset(dataset)
        )

        return waiting_for_move

    def test_dataset_deletion(self):
        """
        A dataset can be deleted, resulting in its removal from the node.
        """
        created = create_dataset(self)

        def delete_dataset(result):
            cluster, dataset = result
            deleted = cluster.delete_dataset(dataset["dataset_id"])

            def not_exists():
                request = cluster.datasets_state()
                request.addCallback(
                    lambda actual_datasets: dataset["dataset_id"] not in
                    (d["dataset_id"] for d in actual_datasets))
                return request
            deleted.addCallback(lambda _: loop_until(not_exists))
            return deleted
        created.addCallback(delete_dataset)
        return created

    def test_dataset_grow(self):
        """
        The size of a dataset can be increased.
        """
        creating = create_dataset(
            test_case=self, maximum_size=REALISTIC_BLOCKDEVICE_SIZE
        )
        new_size = REALISTIC_BLOCKDEVICE_SIZE * 2

        def resize_dataset(result):
            cluster, dataset = result
            return cluster.update_dataset(
                dataset["dataset_id"], {u'maximum_size': new_size}
            )

        resizing = creating.addCallback(resize_dataset)

        def check_dataset_size(result):
            cluster, dataset = result
            self.assertEqual(new_size, dataset['maximum_size'])
            return cluster.wait_for_dataset(dataset)

        checking = resizing.addCallback(check_dataset_size)

        return checking

    def test_dataset_shrink(self):
        """
        The size of a dataset can be decreased.
        """
        creating = create_dataset(
            test_case=self, maximum_size=REALISTIC_BLOCKDEVICE_SIZE * 2
        )
        new_size = REALISTIC_BLOCKDEVICE_SIZE

        def resize_dataset(result):
            cluster, dataset = result
            return cluster.update_dataset(
                dataset["dataset_id"], {u'maximum_size': new_size}
            )

        resizing = creating.addCallback(resize_dataset)

        def check_dataset_size(result):
            cluster, dataset = result
            self.assertEqual(new_size, dataset['maximum_size'])
            return cluster.wait_for_dataset(dataset)

        checking = resizing.addCallback(check_dataset_size)

        return checking

    def test_dataset_shrink_not_valid(self):
        """
        If the requested maximum_size is smaller than the allowed minimum the
        response is ``BAD_REQUEST``.
        """
        creating = create_dataset(
            test_case=self, maximum_size=REALISTIC_BLOCKDEVICE_SIZE
        )
        new_size = 67108864 - 1

        def resize_dataset(result):
            cluster, dataset = result
            # Reconfigure that dataset to be an invalid size.
            resizing = cluster.update_dataset(
                dataset_id=dataset["dataset_id"],
                dataset_properties={u'maximum_size': new_size}
            )
            # Check for expected exception and response code.
            return self.assertFailure(
                resizing, ValueError
            ).addCallback(
                lambda exception: self.assertEqual(
                    BAD_REQUEST, exception.args[1]
                )
            )

        return creating.addCallback(resize_dataset)

    def test_dataset_remove_size_limit(self):
        """
        A dataset with a size limit can have that limit removed.
        """
        creating = create_dataset(
            test_case=self, maximum_size=REALISTIC_BLOCKDEVICE_SIZE
        )
        new_size = None

        def resize_dataset(result):
            cluster, dataset = result
            return cluster.update_dataset(
                dataset["dataset_id"], {u'maximum_size': new_size}
            )

        resizing = creating.addCallback(resize_dataset)

        def check_dataset_size(result):
            cluster, dataset = result
            # If there is no maximum_size, the configuration response will not
            # contain that key
            self.assertNotIn(u'maximum_size', dataset.keys())
            return cluster.wait_for_dataset(dataset)

        checking = resizing.addCallback(check_dataset_size)

        return checking
<|endoftext|>


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

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import collections

# Dependency imports
import six
import tensorflow as tf

from tensorflow_federated.python.common_libs import py_typecheck

from tensorflow_federated.python.core.impl import placement_literals

@six.add_metaclass(abc.ABCMeta)
class Type(object):
  """An abstract interface for all classes that represent TFF types."""

  @abc.abstractmethod
  def is_assignable_from(self, other):
    """Determines whether this TFF type is assignable from another TFF type.

    Args:
      other: Another type, an instance of `Type`.

    Returns:
      `True` if self is assignable from other, `False` otherwise.
    """
    raise NotImplementedError

  @abc.abstractmethod
  def __repr__(self):
    """Returns a full-form representation of this type."""
    raise NotImplementedError

  @abc.abstractmethod
  def __str__(self):
    """Returns a concise representation of this type."""
    raise NotImplementedError

  # Types are partially ordered by the relation of assignability: if type T is
  # assignable from type T', so T' is more narrowly specialized than T, we say
  # that T >= T'.

  def __ge__(self, other):
    return self.is_assignable_from(other)

  def __le__(self, other):
    return other.__ge__(self)

  def __eq__(self, other):
    return self.__ge__(other) and self.__le__(other)

  def __ne__(self, other):
    return not self.__eq__(other)

  def __gt__(self, other):
    return self.__ge__(other) and not self.__le__(other)

  def __lt__(self, other):
    return self.__le__(other) and not self.__ge__(other)

class TensorType(Type):
  """An implementation of `Type` for representing types of tensors in TFF."""

  def __init__(self, dtype, shape=None):
    """Constructs a new instance from the given dtype and shape.

    Args:
      dtype: An instance of `tf.DType`.
      shape: An optional instance of `tf.TensorShape` or an argument that can be
        passed to its constructor (such as a `list` or a `tuple`), or `None` for
        the default scalar shape. Unspecified shapes are not supported.

    Raises:
      TypeError: if arguments are of the wrong types.
    """
    py_typecheck.check_type(dtype, tf.DType)
    self._dtype = dtype
    if shape is None:
      self._shape = tf.TensorShape([])
    elif isinstance(shape, tf.TensorShape):
      self._shape = shape
    else:
      self._shape = tf.TensorShape(shape)

  @property
  def dtype(self):
    return self._dtype

  @property
  def shape(self):
    return self._shape

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)

    def _shape_is_assignable_from(x, y):

      def _dimension_is_assignable_from(x, y):
        # Either the first dimension is undefined or it has the same size as
        # the second.
        return (x.value is None) or (x.value == y.value)

      # Shapes must have equal ranks, and all dimensions in first have to be
      # assignable from the corresponding dimensions in the second.
      return ((x.ndims == y.ndims) and ((x.dims is None) or all(
          _dimension_is_assignable_from(x.dims[k], y.dims[k])
          for k in six.moves.range(x.ndims))))

    return (isinstance(other, TensorType) and (self.dtype == other.dtype) and
            _shape_is_assignable_from(self.shape, other.shape))

  def __repr__(self):
    if self._shape.ndims > 0:
      values = repr([dim.value for dim in self._shape.dims])
      return 'TensorType({}, {})'.format(repr(self._dtype), values)
    else:
      return 'TensorType({})'.format(repr(self._dtype))

  def __str__(self):
    if self._shape.ndims > 0:
      values = [
          str(dim.value) if dim.value is not None else '?'
          for dim in self._shape.dims
      ]
      return '{}[{}]'.format(self._dtype.name, ','.join(values))
    else:
      return self._dtype.name

class NamedTupleType(Type):
  """An implementation of `Type` for representing named tuple types in TFF."""

  def __init__(self, elements):
    """Constructs a new instance from the given element types.

    Args:
      elements: A list of element specifications. Each element specification is
        either a type spec (an instance of `Type` or something convertible to it
        via `to_type()`) for the element, or a pair (name, spec) for elements
        that have defined names. Alternatively, one can supply here an instance
        of `collections.OrderedDict` mapping element names to their types (or
        things that are convertible to types).

    Raises:
      TypeError: if the arguments are of the wrong types.
      ValueError: if the named tuple contains no elements.
    """
    py_typecheck.check_type(elements, (list, tuple, collections.OrderedDict))
    if '_asdict' in vars(type(elements)):
      elements = elements._asdict()
    if isinstance(elements, collections.OrderedDict):
      elements = elements.items()
    if not elements:
      raise ValueError('A named tuple must contain at least one element.')

    def _is_named_element(e):
      return (isinstance(e, tuple) and (len(e) == 2) and
              isinstance(e[0], six.string_types))

    def _map_element(e):
      if isinstance(e, Type):
        return (None, e)
      elif _is_named_element(e):
        return (e[0], to_type(e[1]))
      else:
        return (None, to_type(e))

    if _is_named_element(elements):
      self._elements = [(elements[0], to_type(elements[1]))]
    else:
      self._elements = [_map_element(e) for e in elements]

  @property
  def elements(self):
    # Shallow copy to prevent accidental modification by the caller.
    return list(self._elements)

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)
    if not isinstance(other, NamedTupleType):
      return False
    other_elements = other.elements
    return ((len(self._elements) == len(other_elements)) and all(
        ((self._elements[k][0] in [other_elements[k][0], None]) and
         (self._elements[k][1].is_assignable_from(other_elements[k][1])))
        for k in six.moves.range(len(self._elements))))

  def __repr__(self):

    def _element_repr(e):
      if e[0] is not None:
        return '(\'{}\', {})'.format(e[0], repr(e[1]))
      else:
        return repr(e[1])

    return 'NamedTupleType([{}])'.format(', '.join(
        [_element_repr(e) for e in self._elements]))

  def __str__(self):

    def _element_str(name, value):
      return '{}={}'.format(name, str(value)) if name else str(value)

    return ('<{}>'.format(','.join(
        [_element_str(name, value) for name, value in self._elements])))

class SequenceType(Type):
  """An implementation of `Type` for representing types of sequences in TFF."""

  def __init__(self, element):
    """Constructs a new instance from the given element type.

    Args:
      element: A specification of the element type, either an instance of `Type`
        or something convertible to it by `to_type()`.
    """
    self._element = to_type(element)

  @property
  def element(self):
    return self._element

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)
    return (isinstance(other, SequenceType) and
            self.element.is_assignable_from(other.element))

  def __repr__(self):
    return 'SequenceType({})'.format(repr(self._element))

  def __str__(self):
    return '{}*'.format(str(self._element))

class FunctionType(Type):
  """An implementation of `Type` for representing functional types in TFF."""

  def __init__(self, parameter, result):
    """Constructs a new instance from the given parameter and result types.

    Args:
      parameter: A specification of the parameter type, either an instance of
        `Type` or something convertible to it by `to_type()`.
      result: A specification of the result type, either an instance of `Type`
        or something convertible to it by `to_type()`.
    """
    self._parameter = to_type(parameter)
    self._result = to_type(result)

  @property
  def parameter(self):
    return self._parameter

  @property
  def result(self):
    return self._result

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)
    return (isinstance(other, FunctionType) and
            (((other.parameter is None) and (self.parameter is None)) or
             ((other.parameter is not None) and (self.parameter is not None) and
              other.parameter.is_assignable_from(self.parameter)) and
             self.result.is_assignable_from(other.result)))

  def __repr__(self):
    return 'FunctionType({}, {})'.format(
        repr(self._parameter), repr(self._result))

  def __str__(self):
    return '({} -> {})'.format(
        str(self._parameter) if self._parameter is not None else '',
        str(self._result))

class AbstractType(Type):
  """An implementation of `Type` for representing abstract types in TFF."""

  def __init__(self, label):
    """Constructs a new instance from the given string label.

    Args:
      label: A string label of an abstract type. All occurences of the label
        within a computation's type signature refer to the same concrete type.
    """
    py_typecheck.check_type(label, six.string_types)
    self._label = str(label)

  @property
  def label(self):
    return self._label

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)

    # TODO(b/113112108): Revise this to extend the relation of assignability to
    # abstract types.
    raise ValueError('Abstract types are not comparable.')

  def __repr__(self):
    return 'AbstractType(\'{}\')'.format(self._label)

  def __str__(self):
    return self._label

class PlacementType(Type):
  """An implementation of `Type` for representing the placement type in TFF.

  There is only one placement type, a TFF built-in, just as there is only one
  `int` or `str` type in Python. All instances of this class represent the same
  built-in TFF placement type.
  """

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)
    return isinstance(other, PlacementType)

  def __repr__(self):
    return 'PlacementType()'

  def __str__(self):
    return 'placement'

class FederatedType(Type):
  """An implementation of `Type` for representing federated types in TFF."""

  def __init__(self, member, placement, all_equal=False):
    """Constructs a new federated type instance.

    Args:
      member: An instance of `Type` (or something convertible to it) that
        represents the type of the member components of each value of this
        federated type.
      placement: The specification of placement that the member components of
        this federated type are hosted on. Must be either a placement literal
        such as `SERVER` or `CLIENTS` to refer to a globally defined placement,
        or a placement label to refer to a placement defined in other parts of a
        type signature. Specifying placement labels is not implemented yet.
      all_equal: A `bool` value that indicates whether all members of the
        federated type are equal (`True`), or are allowed to differ (`False`).
    """
    if not isinstance(placement, placement_literals.PlacementLiteral):
      raise NotImplementedError(
          'At the moment, only specifying placement literals is implemented.')
    py_typecheck.check_type(all_equal, bool)
    self._member = to_type(member)
    self._placement = placement
    self._all_equal = all_equal

  # TODO(b/113112108): Extend this to support federated types parameterized
  # by abstract placement labels, such as those used in generic types of
  # federated operators.

  @property
  def member(self):
    return self._member

  @property
  def placement(self):
    return self._placement

  @property
  def all_equal(self):
    return self._all_equal

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)
    if (not isinstance(other, FederatedType) or
        not self._member.is_assignable_from(other.member) or
        self._all_equal and not other.all_equal):
      return False
    for val in [self, other]:
      py_typecheck.check_type(val.placement,
                              placement_literals.PlacementLiteral)
    return self.placement is other.placement

  def __repr__(self):
    return 'FederatedType({}, {}, {})'.format(
        repr(self._member), repr(self._placement), repr(self._all_equal))

  def __str__(self):
    if self._all_equal:
      return '{}@{}'.format(str(self._member), str(self._placement))
    else:
      return '{{{}}}@{}'.format(str(self._member), str(self._placement))

def to_type(spec):
  """Converts the argument into an instance of `Type`.

  Args:
    spec: Either an instance of `Type`, or an argument convertible to Type.
      Assorted examples of type specifications are included below.

      Examples of arguments convertible to tensor types:

        ```
        tf.int32
        (tf.int32, [10])
        (tf.int32, [None])
        ```

      Examples of arguments convertible to flat named tuple types:

        ```
        [tf.int32, tf.bool]
        (tf.int32, tf.bool)
        [('a', tf.int32), ('b', tf.bool)]
        ('a', tf.int32)
        collections.OrderedDict([('a', tf.int32), ('b', tf.bool)])
        ```

      Examples of arguments convertible to nested named tuple types:

        ```
        (tf.int32, (tf.float32, tf.bool))
        (tf.int32, (('x', tf.float32), tf.bool))
        ((tf.int32, [1]), (('x', (tf.float32, [2])), (tf.bool, [3])))
        ```

  Returns:
    An instance of `Type` corresponding to the given spec.
  """
  # TODO(b/113112108): Add multiple examples of valid type specs here in the
  # comments, in addition to the unit test.

  if isinstance(spec, Type) or spec is None:
    return spec
  elif isinstance(spec, tf.DType):
    return TensorType(spec)
  elif (isinstance(spec, tuple) and (len(spec) == 2) and
        isinstance(spec[0], tf.DType) and
        (isinstance(spec[1], tf.TensorShape) or
         (isinstance(spec[1], (list, tuple)) and all(
             (isinstance(x, int) or x is None) for x in spec[1])))):
    # We found a 2-element tuple of the form (dtype, shape), where dtype is an
    # instance of tf.DType, and shape is either an instance of tf.TensorShape,
    # or a list, or a tuple that can be fed as argument into a tf.TensorShape.
    # We thus convert this into a TensorType.
    return TensorType(spec[0], spec[1])
  elif isinstance(spec, (list, tuple, collections.OrderedDict)):
    return NamedTupleType(spec)
  else:
    raise TypeError(
        'Unable to interpret an argument of type {} as a type spec.'.format(
            py_typecheck.type_string(type(spec))))
<commit_msg>Disable formatting of `to_type` docstring.<commit_after># Copyright 2018, The TensorFlow Federated Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Defines functions and classes for building and manipulating TFF types."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import collections

# Dependency imports
import six
import tensorflow as tf

from tensorflow_federated.python.common_libs import py_typecheck

from tensorflow_federated.python.core.impl import placement_literals

@six.add_metaclass(abc.ABCMeta)
class Type(object):
  """An abstract interface for all classes that represent TFF types."""

  @abc.abstractmethod
  def is_assignable_from(self, other):
    """Determines whether this TFF type is assignable from another TFF type.

    Args:
      other: Another type, an instance of `Type`.

    Returns:
      `True` if self is assignable from other, `False` otherwise.
    """
    raise NotImplementedError

  @abc.abstractmethod
  def __repr__(self):
    """Returns a full-form representation of this type."""
    raise NotImplementedError

  @abc.abstractmethod
  def __str__(self):
    """Returns a concise representation of this type."""
    raise NotImplementedError

  # Types are partially ordered by the relation of assignability: if type T is
  # assignable from type T', so T' is more narrowly specialized than T, we say
  # that T >= T'.

  def __ge__(self, other):
    return self.is_assignable_from(other)

  def __le__(self, other):
    return other.__ge__(self)

  def __eq__(self, other):
    return self.__ge__(other) and self.__le__(other)

  def __ne__(self, other):
    return not self.__eq__(other)

  def __gt__(self, other):
    return self.__ge__(other) and not self.__le__(other)

  def __lt__(self, other):
    return self.__le__(other) and not self.__ge__(other)

class TensorType(Type):
  """An implementation of `Type` for representing types of tensors in TFF."""

  def __init__(self, dtype, shape=None):
    """Constructs a new instance from the given dtype and shape.

    Args:
      dtype: An instance of `tf.DType`.
      shape: An optional instance of `tf.TensorShape` or an argument that can be
        passed to its constructor (such as a `list` or a `tuple`), or `None` for
        the default scalar shape. Unspecified shapes are not supported.

    Raises:
      TypeError: if arguments are of the wrong types.
    """
    py_typecheck.check_type(dtype, tf.DType)
    self._dtype = dtype
    if shape is None:
      self._shape = tf.TensorShape([])
    elif isinstance(shape, tf.TensorShape):
      self._shape = shape
    else:
      self._shape = tf.TensorShape(shape)

  @property
  def dtype(self):
    return self._dtype

  @property
  def shape(self):
    return self._shape

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)

    def _shape_is_assignable_from(x, y):

      def _dimension_is_assignable_from(x, y):
        # Either the first dimension is undefined or it has the same size as
        # the second.
        return (x.value is None) or (x.value == y.value)

      # Shapes must have equal ranks, and all dimensions in first have to be
      # assignable from the corresponding dimensions in the second.
      return ((x.ndims == y.ndims) and ((x.dims is None) or all(
          _dimension_is_assignable_from(x.dims[k], y.dims[k])
          for k in six.moves.range(x.ndims))))

    return (isinstance(other, TensorType) and (self.dtype == other.dtype) and
            _shape_is_assignable_from(self.shape, other.shape))

  def __repr__(self):
    if self._shape.ndims > 0:
      values = repr([dim.value for dim in self._shape.dims])
      return 'TensorType({}, {})'.format(repr(self._dtype), values)
    else:
      return 'TensorType({})'.format(repr(self._dtype))

  def __str__(self):
    if self._shape.ndims > 0:
      values = [
          str(dim.value) if dim.value is not None else '?'
          for dim in self._shape.dims
      ]
      return '{}[{}]'.format(self._dtype.name, ','.join(values))
    else:
      return self._dtype.name

class NamedTupleType(Type):
  """An implementation of `Type` for representing named tuple types in TFF."""

  def __init__(self, elements):
    """Constructs a new instance from the given element types.

    Args:
      elements: A list of element specifications. Each element specification is
        either a type spec (an instance of `Type` or something convertible to it
        via `to_type()`) for the element, or a pair (name, spec) for elements
        that have defined names. Alternatively, one can supply here an instance
        of `collections.OrderedDict` mapping element names to their types (or
        things that are convertible to types).

    Raises:
      TypeError: if the arguments are of the wrong types.
      ValueError: if the named tuple contains no elements.
    """
    py_typecheck.check_type(elements, (list, tuple, collections.OrderedDict))
    if '_asdict' in vars(type(elements)):
      elements = elements._asdict()
    if isinstance(elements, collections.OrderedDict):
      elements = elements.items()
    if not elements:
      raise ValueError('A named tuple must contain at least one element.')

    def _is_named_element(e):
      return (isinstance(e, tuple) and (len(e) == 2) and
              isinstance(e[0], six.string_types))

    def _map_element(e):
      if isinstance(e, Type):
        return (None, e)
      elif _is_named_element(e):
        return (e[0], to_type(e[1]))
      else:
        return (None, to_type(e))

    if _is_named_element(elements):
      self._elements = [(elements[0], to_type(elements[1]))]
    else:
      self._elements = [_map_element(e) for e in elements]

  @property
  def elements(self):
    # Shallow copy to prevent accidental modification by the caller.
    return list(self._elements)

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)
    if not isinstance(other, NamedTupleType):
      return False
    other_elements = other.elements
    return ((len(self._elements) == len(other_elements)) and all(
        ((self._elements[k][0] in [other_elements[k][0], None]) and
         (self._elements[k][1].is_assignable_from(other_elements[k][1])))
        for k in six.moves.range(len(self._elements))))

  def __repr__(self):

    def _element_repr(e):
      if e[0] is not None:
        return '(\'{}\', {})'.format(e[0], repr(e[1]))
      else:
        return repr(e[1])

    return 'NamedTupleType([{}])'.format(', '.join(
        [_element_repr(e) for e in self._elements]))

  def __str__(self):

    def _element_str(name, value):
      return '{}={}'.format(name, str(value)) if name else str(value)

    return ('<{}>'.format(','.join(
        [_element_str(name, value) for name, value in self._elements])))

class SequenceType(Type):
  """An implementation of `Type` for representing types of sequences in TFF."""

  def __init__(self, element):
    """Constructs a new instance from the given element type.

    Args:
      element: A specification of the element type, either an instance of `Type`
        or something convertible to it by `to_type()`.
    """
    self._element = to_type(element)

  @property
  def element(self):
    return self._element

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)
    return (isinstance(other, SequenceType) and
            self.element.is_assignable_from(other.element))

  def __repr__(self):
    return 'SequenceType({})'.format(repr(self._element))

  def __str__(self):
    return '{}*'.format(str(self._element))

class FunctionType(Type):
  """An implementation of `Type` for representing functional types in TFF."""

  def __init__(self, parameter, result):
    """Constructs a new instance from the given parameter and result types.

    Args:
      parameter: A specification of the parameter type, either an instance of
        `Type` or something convertible to it by `to_type()`.
      result: A specification of the result type, either an instance of `Type`
        or something convertible to it by `to_type()`.
    """
    self._parameter = to_type(parameter)
    self._result = to_type(result)

  @property
  def parameter(self):
    return self._parameter

  @property
  def result(self):
    return self._result

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)
    return (isinstance(other, FunctionType) and
            (((other.parameter is None) and (self.parameter is None)) or
             ((other.parameter is not None) and (self.parameter is not None) and
              other.parameter.is_assignable_from(self.parameter)) and
             self.result.is_assignable_from(other.result)))

  def __repr__(self):
    return 'FunctionType({}, {})'.format(
        repr(self._parameter), repr(self._result))

  def __str__(self):
    return '({} -> {})'.format(
        str(self._parameter) if self._parameter is not None else '',
        str(self._result))

class AbstractType(Type):
  """An implementation of `Type` for representing abstract types in TFF."""

  def __init__(self, label):
    """Constructs a new instance from the given string label.

    Args:
      label: A string label of an abstract type. All occurences of the label
        within a computation's type signature refer to the same concrete type.
    """
    py_typecheck.check_type(label, six.string_types)
    self._label = str(label)

  @property
  def label(self):
    return self._label

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)

    # TODO(b/113112108): Revise this to extend the relation of assignability to
    # abstract types.
    raise ValueError('Abstract types are not comparable.')

  def __repr__(self):
    return 'AbstractType(\'{}\')'.format(self._label)

  def __str__(self):
    return self._label

class PlacementType(Type):
  """An implementation of `Type` for representing the placement type in TFF.

  There is only one placement type, a TFF built-in, just as there is only one
  `int` or `str` type in Python. All instances of this class represent the same
  built-in TFF placement type.
  """

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)
    return isinstance(other, PlacementType)

  def __repr__(self):
    return 'PlacementType()'

  def __str__(self):
    return 'placement'

class FederatedType(Type):
  """An implementation of `Type` for representing federated types in TFF."""

  def __init__(self, member, placement, all_equal=False):
    """Constructs a new federated type instance.

    Args:
      member: An instance of `Type` (or something convertible to it) that
        represents the type of the member components of each value of this
        federated type.
      placement: The specification of placement that the member components of
        this federated type are hosted on. Must be either a placement literal
        such as `SERVER` or `CLIENTS` to refer to a globally defined placement,
        or a placement label to refer to a placement defined in other parts of a
        type signature. Specifying placement labels is not implemented yet.
      all_equal: A `bool` value that indicates whether all members of the
        federated type are equal (`True`), or are allowed to differ (`False`).
    """
    if not isinstance(placement, placement_literals.PlacementLiteral):
      raise NotImplementedError(
          'At the moment, only specifying placement literals is implemented.')
    py_typecheck.check_type(all_equal, bool)
    self._member = to_type(member)
    self._placement = placement
    self._all_equal = all_equal

  # TODO(b/113112108): Extend this to support federated types parameterized
  # by abstract placement labels, such as those used in generic types of
  # federated operators.

  @property
  def member(self):
    return self._member

  @property
  def placement(self):
    return self._placement

  @property
  def all_equal(self):
    return self._all_equal

  def is_assignable_from(self, other):
    py_typecheck.check_type(other, Type)
    if (not isinstance(other, FederatedType) or
        not self._member.is_assignable_from(other.member) or
        self._all_equal and not other.all_equal):
      return False
    for val in [self, other]:
      py_typecheck.check_type(val.placement,
                              placement_literals.PlacementLiteral)
    return self.placement is other.placement

  def __repr__(self):
    return 'FederatedType({}, {}, {})'.format(
        repr(self._member), repr(self._placement), repr(self._all_equal))

  def __str__(self):
    if self._all_equal:
      return '{}@{}'.format(str(self._member), str(self._placement))
    else:
      return '{{{}}}@{}'.format(str(self._member), str(self._placement))

def to_type(spec):
  # pyformat: disable
  """Converts the argument into an instance of `Type`.

  Args:
    spec: Either an instance of `Type`, or an argument convertible to Type.
      Assorted examples of type specifications are included below.

      Examples of arguments convertible to tensor types:

      ```
      tf.int32
      (tf.int32, [10])
      (tf.int32, [None])
      ```

      Examples of arguments convertible to flat named tuple types:

      ```
      [tf.int32, tf.bool]
      (tf.int32, tf.bool)
      [('a', tf.int32), ('b', tf.bool)]
      ('a', tf.int32)
      collections.OrderedDict([('a', tf.int32), ('b', tf.bool)])
      ```

      Examples of arguments convertible to nested named tuple types:

      ```
      (tf.int32, (tf.float32, tf.bool))
      (tf.int32, (('x', tf.float32), tf.bool))
      ((tf.int32, [1]), (('x', (tf.float32, [2])), (tf.bool, [3])))
      ```

  Returns:
    An instance of `Type` corresponding to the given spec.
  """
  # pyformat: enable
  # TODO(b/113112108): Add multiple examples of valid type specs here in the
  # comments, in addition to the unit test.

  if isinstance(spec, Type) or spec is None:
    return spec
  elif isinstance(spec, tf.DType):
    return TensorType(spec)
  elif (isinstance(spec, tuple) and (len(spec) == 2) and
        isinstance(spec[0], tf.DType) and
        (isinstance(spec[1], tf.TensorShape) or
         (isinstance(spec[1], (list, tuple)) and all(
             (isinstance(x, int) or x is None) for x in spec[1])))):
    # We found a 2-element tuple of the form (dtype, shape), where dtype is an
    # instance of tf.DType, and shape is either an instance of tf.TensorShape,
    # or a list, or a tuple that can be fed as argument into a tf.TensorShape.
    # We thus convert this into a TensorType.
    return TensorType(spec[0], spec[1])
  elif isinstance(spec, (list, tuple, collections.OrderedDict)):
    return NamedTupleType(spec)
  else:
    raise TypeError(
        'Unable to interpret an argument of type {} as a type spec.'.format(
            py_typecheck.type_string(type(spec))))
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# Copyright (c) 2014-2016 The Bitcoin Core developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
"""Test block proposals with getblocktemplate."""

from test_framework.test_framework import BitcoinTestFramework
from test_framework.util import *

from binascii import a2b_hex, b2a_hex
from hashlib import sha256
from struct import pack

def b2x(b):
    return b2a_hex(b).decode('ascii')

# NOTE: This does not work for signed numbers (set the high bit) or zero (use b'\0')
def encodeUNum(n):
    s = bytearray(b'\1')
    while n > 127:
        s[0] += 1
        s.append(n % 256)
        n //= 256
    s.append(n)
    return bytes(s)

def varlenEncode(n):
    if n < 0xfd:
        return pack('<B', n)
    if n <= 0xffff:
        return b'\xfd' + pack('<H', n)
    if n <= 0xffffffff:
        return b'\xfe' + pack('<L', n)
    return b'\xff' + pack('<Q', n)

def dblsha(b):
    return sha256(sha256(b).digest()).digest()

def genmrklroot(leaflist):
    cur = leaflist
    while len(cur) > 1:
        n = []
        if len(cur) & 1:
            cur.append(cur[-1])
        for i in range(0, len(cur), 2):
            n.append(dblsha(cur[i] + cur[i+1]))
        cur = n
    return cur[0]

def template_to_bytearray(tmpl, txlist):
    blkver = pack('<L', tmpl['version'])
    mrklroot = genmrklroot(list(dblsha(a) for a in txlist))
    timestamp = pack('<L', tmpl['curtime'])
    nonce = b'\0\0\0\0'
    blk = blkver + a2b_hex(tmpl['previousblockhash'])[::-1] + mrklroot + timestamp + a2b_hex(tmpl['bits'])[::-1] + nonce
    blk += varlenEncode(len(txlist))
    for tx in txlist:
        blk += tx
    return bytearray(blk)

def template_to_hex(tmpl, txlist):
    return b2x(template_to_bytearray(tmpl, txlist))

def assert_template(node, tmpl, txlist, expect):
    rsp = node.getblocktemplate({'data':template_to_hex(tmpl, txlist),'mode':'proposal'})
    if rsp != expect:
        raise AssertionError('unexpected: %s' % (rsp,))

class GetBlockTemplateProposalTest(BitcoinTestFramework):

    def __init__(self):
        super().__init__()
        self.num_nodes = 2
        self.setup_clean_chain = False

    def run_test(self):
        node = self.nodes[0]
        node.generate(1) # Mine a block to leave initial block download
        tmpl = node.getblocktemplate()
        if 'coinbasetxn' not in tmpl:
            rawcoinbase = encodeUNum(tmpl['height'])
            rawcoinbase += b'\x01-'
            hexcoinbase = b2x(rawcoinbase)
            hexoutval = b2x(pack('<Q', tmpl['coinbasevalue']))
            tmpl['coinbasetxn'] = {'data': '01000000' + '01' + '0000000000000000000000000000000000000000000000000000000000000000ffffffff' + ('%02x' % (len(rawcoinbase),)) + hexcoinbase + 'fffffffe' + '01' + hexoutval + '00' + '00000000'}
        txlist = list(bytearray(a2b_hex(a['data'])) for a in (tmpl['coinbasetxn'],) + tuple(tmpl['transactions']))

        # Test 0: Capability advertised
        assert('proposal' in tmpl['capabilities'])

        # NOTE: This test currently FAILS (regtest mode doesn't enforce block height in coinbase)
        ## Test 1: Bad height in coinbase
        #txlist[0][4+1+36+1+1] += 1
        #assert_template(node, tmpl, txlist, 'FIXME')
        #txlist[0][4+1+36+1+1] -= 1

        # Test 2: Bad input hash for gen tx
        txlist[0][4+1] += 1
        assert_template(node, tmpl, txlist, 'bad-cb-missing')
        txlist[0][4+1] -= 1

        # Test 3: Truncated final tx
        lastbyte = txlist[-1].pop()
        assert_raises_jsonrpc(-22, "Block decode failed", assert_template, node, tmpl, txlist, 'n/a')
        txlist[-1].append(lastbyte)

        # Test 4: Add an invalid tx to the end (duplicate of gen tx)
        txlist.append(txlist[0])
        assert_template(node, tmpl, txlist, 'bad-txns-duplicate')
        txlist.pop()

        # Test 5: Add an invalid tx to the end (non-duplicate)
        txlist.append(bytearray(txlist[0]))
        txlist[-1][4+1] = 0xff
        assert_template(node, tmpl, txlist, 'bad-txns-inputs-missingorspent')
        txlist.pop()

        # Test 6: Future tx lock time
        txlist[0][-4:] = b'\xff\xff\xff\xff'
        assert_template(node, tmpl, txlist, 'bad-txns-nonfinal')
        txlist[0][-4:] = b'\0\0\0\0'

        # Test 7: Bad tx count
        txlist.append(b'')
        assert_raises_jsonrpc(-22, 'Block decode failed', assert_template, node, tmpl, txlist, 'n/a')
        txlist.pop()

        # Test 8: Bad bits
        realbits = tmpl['bits']
        tmpl['bits'] = '1c0000ff'  # impossible in the real world
        assert_template(node, tmpl, txlist, 'bad-diffbits')
        tmpl['bits'] = realbits

        # Test 9: Bad merkle root
        rawtmpl = template_to_bytearray(tmpl, txlist)
        rawtmpl[4+32] = (rawtmpl[4+32] + 1) % 0x100
        rsp = node.getblocktemplate({'data':b2x(rawtmpl),'mode':'proposal'})
        if rsp != 'bad-txnmrklroot':
            raise AssertionError('unexpected: %s' % (rsp,))

        # Test 10: Bad timestamps
        realtime = tmpl['curtime']
        tmpl['curtime'] = 0x7fffffff
        assert_template(node, tmpl, txlist, 'time-too-new')
        tmpl['curtime'] = 0
        assert_template(node, tmpl, txlist, 'time-too-old')
        tmpl['curtime'] = realtime

        # Test 11: Valid block
        assert_template(node, tmpl, txlist, None)

        # Test 12: Orphan block
        tmpl['previousblockhash'] = 'ff00' * 16
        assert_template(node, tmpl, txlist, 'inconclusive-not-best-prevblk')

if __name__ == '__main__':
    GetBlockTemplateProposalTest().main()
<commit_msg>[tests] Fix import order in getblocktemplate test<commit_after>#!/usr/bin/env python3
# Copyright (c) 2014-2016 The Bitcoin Core developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
"""Test block proposals with getblocktemplate."""

from binascii import a2b_hex, b2a_hex
from hashlib import sha256
from struct import pack

from test_framework.test_framework import BitcoinTestFramework
from test_framework.util import *

def b2x(b):
    return b2a_hex(b).decode('ascii')

# NOTE: This does not work for signed numbers (set the high bit) or zero (use b'\0')
def encodeUNum(n):
    s = bytearray(b'\1')
    while n > 127:
        s[0] += 1
        s.append(n % 256)
        n //= 256
    s.append(n)
    return bytes(s)

def varlenEncode(n):
    if n < 0xfd:
        return pack('<B', n)
    if n <= 0xffff:
        return b'\xfd' + pack('<H', n)
    if n <= 0xffffffff:
        return b'\xfe' + pack('<L', n)
    return b'\xff' + pack('<Q', n)

def dblsha(b):
    return sha256(sha256(b).digest()).digest()

def genmrklroot(leaflist):
    cur = leaflist
    while len(cur) > 1:
        n = []
        if len(cur) & 1:
            cur.append(cur[-1])
        for i in range(0, len(cur), 2):
            n.append(dblsha(cur[i] + cur[i+1]))
        cur = n
    return cur[0]

def template_to_bytearray(tmpl, txlist):
    blkver = pack('<L', tmpl['version'])
    mrklroot = genmrklroot(list(dblsha(a) for a in txlist))
    timestamp = pack('<L', tmpl['curtime'])
    nonce = b'\0\0\0\0'
    blk = blkver + a2b_hex(tmpl['previousblockhash'])[::-1] + mrklroot + timestamp + a2b_hex(tmpl['bits'])[::-1] + nonce
    blk += varlenEncode(len(txlist))
    for tx in txlist:
        blk += tx
    return bytearray(blk)

def template_to_hex(tmpl, txlist):
    return b2x(template_to_bytearray(tmpl, txlist))

def assert_template(node, tmpl, txlist, expect):
    rsp = node.getblocktemplate({'data':template_to_hex(tmpl, txlist),'mode':'proposal'})
    if rsp != expect:
        raise AssertionError('unexpected: %s' % (rsp,))

class GetBlockTemplateProposalTest(BitcoinTestFramework):

    def __init__(self):
        super().__init__()
        self.num_nodes = 2
        self.setup_clean_chain = False

    def run_test(self):
        node = self.nodes[0]
        node.generate(1) # Mine a block to leave initial block download
        tmpl = node.getblocktemplate()
        if 'coinbasetxn' not in tmpl:
            rawcoinbase = encodeUNum(tmpl['height'])
            rawcoinbase += b'\x01-'
            hexcoinbase = b2x(rawcoinbase)
            hexoutval = b2x(pack('<Q', tmpl['coinbasevalue']))
            tmpl['coinbasetxn'] = {'data': '01000000' + '01' + '0000000000000000000000000000000000000000000000000000000000000000ffffffff' + ('%02x' % (len(rawcoinbase),)) + hexcoinbase + 'fffffffe' + '01' + hexoutval + '00' + '00000000'}
        txlist = list(bytearray(a2b_hex(a['data'])) for a in (tmpl['coinbasetxn'],) + tuple(tmpl['transactions']))

        # Test 0: Capability advertised
        assert('proposal' in tmpl['capabilities'])

        # NOTE: This test currently FAILS (regtest mode doesn't enforce block height in coinbase)
        ## Test 1: Bad height in coinbase
        #txlist[0][4+1+36+1+1] += 1
        #assert_template(node, tmpl, txlist, 'FIXME')
        #txlist[0][4+1+36+1+1] -= 1

        # Test 2: Bad input hash for gen tx
        txlist[0][4+1] += 1
        assert_template(node, tmpl, txlist, 'bad-cb-missing')
        txlist[0][4+1] -= 1

        # Test 3: Truncated final tx
        lastbyte = txlist[-1].pop()
        assert_raises_jsonrpc(-22, "Block decode failed", assert_template, node, tmpl, txlist, 'n/a')
        txlist[-1].append(lastbyte)

        # Test 4: Add an invalid tx to the end (duplicate of gen tx)
        txlist.append(txlist[0])
        assert_template(node, tmpl, txlist, 'bad-txns-duplicate')
        txlist.pop()

        # Test 5: Add an invalid tx to the end (non-duplicate)
        txlist.append(bytearray(txlist[0]))
        txlist[-1][4+1] = 0xff
        assert_template(node, tmpl, txlist, 'bad-txns-inputs-missingorspent')
        txlist.pop()

        # Test 6: Future tx lock time
        txlist[0][-4:] = b'\xff\xff\xff\xff'
        assert_template(node, tmpl, txlist, 'bad-txns-nonfinal')
        txlist[0][-4:] = b'\0\0\0\0'

        # Test 7: Bad tx count
        txlist.append(b'')
        assert_raises_jsonrpc(-22, 'Block decode failed', assert_template, node, tmpl, txlist, 'n/a')
        txlist.pop()

        # Test 8: Bad bits
        realbits = tmpl['bits']
        tmpl['bits'] = '1c0000ff'  # impossible in the real world
        assert_template(node, tmpl, txlist, 'bad-diffbits')
        tmpl['bits'] = realbits

        # Test 9: Bad merkle root
        rawtmpl = template_to_bytearray(tmpl, txlist)
        rawtmpl[4+32] = (rawtmpl[4+32] + 1) % 0x100
        rsp = node.getblocktemplate({'data':b2x(rawtmpl),'mode':'proposal'})
        if rsp != 'bad-txnmrklroot':
            raise AssertionError('unexpected: %s' % (rsp,))

        # Test 10: Bad timestamps
        realtime = tmpl['curtime']
        tmpl['curtime'] = 0x7fffffff
        assert_template(node, tmpl, txlist, 'time-too-new')
        tmpl['curtime'] = 0
        assert_template(node, tmpl, txlist, 'time-too-old')
        tmpl['curtime'] = realtime

        # Test 11: Valid block
        assert_template(node, tmpl, txlist, None)

        # Test 12: Orphan block
        tmpl['previousblockhash'] = 'ff00' * 16
        assert_template(node, tmpl, txlist, 'inconclusive-not-best-prevblk')

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


<commit_before># Copyright (C) 2016 Google Inc.
# Licensed under http://www.apache.org/licenses/LICENSE-2.0 <see LICENSE file>

"""Tests for snapshot model."""

from ggrc.app import app
from ggrc.models import all_models
from integration.ggrc.converters import TestCase
from integration.ggrc.models import factories

class TestSnapshot(TestCase):
  """Basic tests for /query api."""

  def setUp(self):
    """Set up test cases for all tests."""
    TestCase.clear_data()
    self._create_cas()
    self._import_file("all_snapshottable_objects.csv")

  @staticmethod
  def _create_cas():
    """Create custom attribute definitions."""
    for type_ in ["facility", "control", "market", "section"]:
      with app.app_context():
        factories.CustomAttributeDefinitionFactory(
            title="CA dropdown",
            definition_type=type_,
            attribute_type="Dropdown",
            multi_choice_options="one,two,three,four,five",
        )
        factories.CustomAttributeDefinitionFactory(
            title="CA text",
            definition_type=type_,
            attribute_type="Text",
        )
        factories.CustomAttributeDefinitionFactory(
            title="CA date",
            definition_type=type_,
            attribute_type="Date",
        )

  def test_revision_conent(self):
    """Test that revision contains all content needed."""

    facility_revision = all_models.Revision.query.filter(
        all_models.Revision.resource_type == "Facility").order_by(
        all_models.Revision.id.desc()).first()

    self.assertIn("custom_attributes", facility_revision.content)
    self.assertNotEqual(facility_revision.content["custom_attributes"], [])
<commit_msg>Update snapshots tests<commit_after># Copyright (C) 2016 Google Inc.
# Licensed under http://www.apache.org/licenses/LICENSE-2.0 <see LICENSE file>

"""Tests for snapshot model."""

from ggrc import db
from ggrc.app import app
from ggrc.models import all_models
from integration.ggrc.converters import TestCase
from integration.ggrc.models import factories

TEST_MODELS = [
    all_models.Facility,
    all_models.AccessGroup,
    all_models.Clause,
    all_models.Contract,
    all_models.Control,
    all_models.DataAsset,
    all_models.Market,
    all_models.Objective,
    all_models.OrgGroup,
    all_models.Policy,
    all_models.Process,
    all_models.Product,
    all_models.Regulation,
    all_models.Risk,
    all_models.Section,
    all_models.Standard,
    all_models.System,
    all_models.Threat,
    all_models.Vendor,
]

class TestSnapshot(TestCase):
  """Basic tests for /query api."""

  IGNORE_KEYS = {
    # currently not working fields:
    "audit_duration",
    "audit_duration_id",

    "audit_frequency",
    "audit_frequency_id",

    "directive",
    "directive_id",

    "kind",
    "kind_id",

    "means",
    "means_id",

    "meta_kind",

    "network_zone",
    "network_zone_id",

    "verify_frequency",
    "verify_frequency_id",

    "assertions",
    "categories",
    "categorizations",
    "categorized_assertions",

    # special fields not needed for snapshots.
    "display_name",
    "preconditions_failed",
    "type",
    "workflow_state",

    "selfLink",
    "viewLink",

    # relationships and mappings
    "audit_objects",
    "audits",
    "controls",
    "object_owners",
    "object_people",
    "objects",
    "people",
    "related_destinations",
    "related_sources",
    "risk_objects",
    "risks",
    "task_group_objects",
    "task_group_tasks",
    "task_groups",

    "children",
    "parent",
    "parent_id",

    # we don't need context for snapshots since they are all under an audit.
    "context",
    "context_id",

    # obsolete fields that will be removed
    "custom_attributes",

    # following fields have been handled in fields without _id prefix. That
    # means that "contact" fields should exists and have correct values.
    "contact_id",
    "secondary_contact_id",

    "principal_assessor_id",
    "secondary_assessor_id",

    "modified_by_id",

    "attribute_object",
    "attribute_object_id",
    "attribute_value",
  }

  def setUp(self):
    """Set up test cases for all tests."""
    TestCase.clear_data()
    self._create_cas()
    self._import_file("all_snapshottable_objects.csv")

  @staticmethod
  def _create_cas():
    """Create custom attribute definitions."""
    for type_ in ["facility", "control", "market", "section"]:
      with app.app_context():
        factories.CustomAttributeDefinitionFactory(
            title="CA dropdown",
            definition_type=type_,
            attribute_type="Dropdown",
            multi_choice_options="one,two,three,four,five",
        )
        factories.CustomAttributeDefinitionFactory(
            title="CA text",
            definition_type=type_,
            attribute_type="Text",
        )
        factories.CustomAttributeDefinitionFactory(
            title="CA date",
            definition_type=type_,
            attribute_type="Date",
        )

  def test_revision_conent(self):
    """Test that revision contains all content needed."""

    facility_revision = all_models.Revision.query.filter(
        all_models.Revision.resource_type == "Facility").order_by(
        all_models.Revision.id.desc()).first()

    self.assertIn("custom_attribute_values", facility_revision.content)
    self.assertNotEqual(facility_revision.content["custom_attribute_values"], [])

  def _get_object(self, obj):
    return self.client.get(
        "/api/{}/{}".format(obj._inflector.table_plural, obj.id)
    ).json[obj._inflector.table_singular]

  def _clean_json(self, content):
    """Remove ignored items from JSON content.

    This function removes all ignored items from dicts, changes dates to
    isoformat changes values to int or unicode, so that the end result is a
    dict that can be compared with the JSON dict that was received from the
    server.

    Args:
      content: object that we want to clean, it can be a dict list or a value.

    Returns:
      content with all values cleaned up
    """
    if isinstance(content, list):
      return sorted(self._clean_json(value) for value in content)

    if hasattr(content, 'isoformat'):
      return unicode(content.isoformat())

    if isinstance(content, int):
      # We convert all numbers to the same type so that the diff of a failed
      # test looks nicer. This conversion does not affect the test results just
      # the output.
      return long(content)

    if not isinstance(content, dict):
      return content

    clean = {}
    for key, value in content.items():
      if key not in self.IGNORE_KEYS:
        clean[str(key)] = self._clean_json(value)

    return clean

  def test_snapshot_content(self):
    snapshot = all_models.Snapshot()
    self.client.get("/login")
    for model in TEST_MODELS:
      obj = model.eager_query().first()
      generated_json = self._clean_json(obj.log_json())
      expected_json = self._clean_json(self._get_object(obj))
      self.assertEqual(expected_json, generated_json)
<|endoftext|>


<commit_before><commit_msg>Update anymail backend<commit_after><|endoftext|>


<commit_before>from django.db import models
from django.contrib.contenttypes.models import ContentType
from django.contrib.contenttypes import generic

from datetime import datetime

from djapian import utils

class ChangeManager(models.Manager):
    def create(self, object, action, **kwargs):
        ct = ContentType.objects.get_for_model(object.__class__)

        try:
            old_change = self.get(
                content_type=ct,
                object_id=object.pk
            )
            if old_change.action=="add":
                if action=="edit":
                    old_change.save()
                    return old_change
                elif action=="delete":
                    old_change.delete()
                    return None
            old_change.delete()
        except self.model.DoesNotExist:
            old_change = self.model(object=object)

        old_change.action = action
        old_change.save()

        return old_change

class Change(models.Model):
    ACTIOINS = (
        ("add", "object added"),
        ("edit", "object edited"),
        ("delete", "object deleted"),
    )

    content_type = models.ForeignKey(ContentType, db_index=True)
    object_id = models.PositiveIntegerField()
    date = models.DateTimeField(default=datetime.now)
    action = models.CharField(max_length=6, choices=ACTIOINS)

    object = generic.GenericForeignKey()

    objects = ChangeManager()

    def __unicode__(self):
        return u'%s#%d To action:`%s`, added on %s' % (
            self.content_type,
            self.object_id,
            self.action,
            self.date
        )

    def save(self):
        self.date = datetime.now()

        super(Change, self).save()

    class Meta:
        unique_together = [("content_type", "object_id")]
<commit_msg>fixed #58 with object_id type issue<commit_after>from django.db import models
from django.contrib.contenttypes.models import ContentType
from django.contrib.contenttypes import generic

from datetime import datetime

from djapian import utils

class ChangeManager(models.Manager):
    def create(self, object, action, **kwargs):
        ct = ContentType.objects.get_for_model(object.__class__)

        try:
            old_change = self.get(
                content_type=ct,
                object_id=object.pk
            )
            if old_change.action=="add":
                if action=="edit":
                    old_change.save()
                    return old_change
                elif action=="delete":
                    old_change.delete()
                    return None
            old_change.delete()
        except self.model.DoesNotExist:
            old_change = self.model(object=object)

        old_change.action = action
        old_change.save()

        return old_change

class Change(models.Model):
    ACTIONS = (
        ("add", "object added"),
        ("edit", "object edited"),
        ("delete", "object deleted"),
    )

    content_type = models.ForeignKey(ContentType, db_index=True)
    object_id = models.CharField(max_length=150)

    date = models.DateTimeField(default=datetime.now)
    action = models.CharField(max_length=6, choices=ACTIONS)

    object = generic.GenericForeignKey()

    objects = ChangeManager()

    def __unicode__(self):
        return u'%s#%d To action:`%s`, added on %s' % (
            self.content_type,
            self.object_id,
            self.action,
            self.date
        )

    def save(self):
        self.date = datetime.now()

        super(Change, self).save()

    class Meta:
        unique_together = [("content_type", "object_id")]
<|endoftext|>


<commit_before>import time
from distutils.version import LooseVersion
from functools import wraps
from unittest import TestCase

from nose import SkipTest
import elasticsearch

from elasticutils import get_es, S

class ESTestCase(TestCase):
    """Superclass for Elasticsearch-using test cases.

    :property index_name: name of the index to use
    :property mapping_type_name: the mapping type name
    :property es_settings: settings to use to build a elasticsearch
        Elasticsearch object.
    :property mapping: the mapping to use when creating an index
    :property data: any data to add to the index in setup_class
    :property skip_tests: if Elasticsearch isn't available, then this
        is True and therefore tests should be skipped for this class

    For examples of usage, see the other ``test_*.py`` files.

    """
    index_name = 'elasticutilstest'
    mapping_type_name = 'elasticutilsmappingtype'
    es_settings = {
        'urls': ['http://localhost:9200']
        }
    mapping = {}
    data = []
    skip_tests = False

    @classmethod
    def setup_class(cls):
        """Class setup for tests.

        Checks to see if ES is running and if not, sets ``skip_test``
        to True on the class.
        """
        # Note: TestCase has no setup_class
        try:
            get_es().cluster.health()
        except elasticsearch.TransportError:
            cls.skip_tests = True
            return

        if cls.data:
            cls.create_index(settings={'mappings': cls.mapping})
            cls.index_data(cls.data)
            cls.refresh()

    @classmethod
    def teardown_class(cls):
        """Class tear down for tests."""
        if cls.skip_tests:
            return

        cls.cleanup_index()

    def setUp(self):
        """Set up a single test.

        :raises SkipTest: if ``skip_tests`` is True for this
            class/instance
        """
        if self.skip_tests:
            raise SkipTest
        super(ESTestCase, self).setUp()

    @classmethod
    def get_es(cls):
        return get_es(**cls.es_settings)

    @classmethod
    def get_s(cls, mapping_type=None):
        if mapping_type is not None:
            s = S(mapping_type)
        else:
            s = S()
        return (s.es(**cls.es_settings)
                 .indexes(cls.index_name)
                 .doctypes(cls.mapping_type_name))

    @classmethod
    def create_index(cls, settings=None):
        es = cls.get_es()
        try:
            es.indices.delete(cls.index_name)
        except elasticsearch.NotFoundError:
            pass
        body = {}
        if settings:
            body['settings'] = settings
        es.indices.create(cls.index_name, body=body)

    @classmethod
    def index_data(cls, data, index=None, doctype=None):
        index = index or cls.index_name
        doctype = doctype or cls.mapping_type_name

        es = cls.get_es()

        # TODO: change this to a bulk index
        for item in data:
            es.index(index, doctype, item, id=item['id'])

        cls.refresh()

    @classmethod
    def cleanup_index(cls):
        es = cls.get_es()
        try:
            es.indices.delete(cls.index_name)
        except elasticsearch.NotFoundError:
            pass

    @classmethod
    def refresh(cls):
        """Refresh index after indexing.

        This refreshes the index specified by `self.index_name`.

        """
        cls.get_es().indices.refresh(cls.index_name)
        cls.get_es().cluster.health(wait_for_status='yellow')

def facet_counts_dict(qs, field):
    return dict((t['term'], t['count']) for t in qs.facet_counts()[field])

def require_version(minimum_version):
    """Skip the test if the Elasticsearch version is less than specified.

    :arg minimum_version: string; the minimum Elasticsearch version required

    """

    def decorated(test):
        """Decorator to only run the test if ES version is greater or
        equal than specified.

        """

        @wraps(test)
        def test_with_version(self):
            "Only run the test if ES version is not less than specified."
            actual_version = self.get_es().info()['version']['number']

            if LooseVersion(actual_version) >= LooseVersion(minimum_version):
                test(self)
            else:
                raise SkipTest

        return test_with_version

    return decorated
<commit_msg>Fix tests to not use docstring as name<commit_after>import time
from distutils.version import LooseVersion
from functools import wraps
from unittest import TestCase

from nose import SkipTest
import elasticsearch

from elasticutils import get_es, S

class ESTestCase(TestCase):
    """Superclass for Elasticsearch-using test cases.

    :property index_name: name of the index to use
    :property mapping_type_name: the mapping type name
    :property es_settings: settings to use to build a elasticsearch
        Elasticsearch object.
    :property mapping: the mapping to use when creating an index
    :property data: any data to add to the index in setup_class
    :property skip_tests: if Elasticsearch isn't available, then this
        is True and therefore tests should be skipped for this class

    For examples of usage, see the other ``test_*.py`` files.

    """
    index_name = 'elasticutilstest'
    mapping_type_name = 'elasticutilsmappingtype'
    es_settings = {
        'urls': ['http://localhost:9200']
        }
    mapping = {}
    data = []
    skip_tests = False

    @classmethod
    def setup_class(cls):
        """Class setup for tests.

        Checks to see if ES is running and if not, sets ``skip_test``
        to True on the class.
        """
        # Note: TestCase has no setup_class
        try:
            get_es().cluster.health()
        except elasticsearch.TransportError:
            cls.skip_tests = True
            return

        if cls.data:
            cls.create_index(settings={'mappings': cls.mapping})
            cls.index_data(cls.data)
            cls.refresh()

    @classmethod
    def teardown_class(cls):
        """Class tear down for tests."""
        if cls.skip_tests:
            return

        cls.cleanup_index()

    def setUp(self):
        """Set up a single test.

        :raises SkipTest: if ``skip_tests`` is True for this
            class/instance
        """
        if self.skip_tests:
            raise SkipTest
        super(ESTestCase, self).setUp()

    def shortDescription(self):
        # Prevent the docstring being used as the test name because
        # that's irritating as all hell when trying to fix tests.
        pass

    @classmethod
    def get_es(cls):
        return get_es(**cls.es_settings)

    @classmethod
    def get_s(cls, mapping_type=None):
        if mapping_type is not None:
            s = S(mapping_type)
        else:
            s = S()
        return (s.es(**cls.es_settings)
                 .indexes(cls.index_name)
                 .doctypes(cls.mapping_type_name))

    @classmethod
    def create_index(cls, settings=None):
        es = cls.get_es()
        try:
            es.indices.delete(cls.index_name)
        except elasticsearch.NotFoundError:
            pass
        body = {}
        if settings:
            body['settings'] = settings
        es.indices.create(cls.index_name, body=body)

    @classmethod
    def index_data(cls, data, index=None, doctype=None):
        index = index or cls.index_name
        doctype = doctype or cls.mapping_type_name

        es = cls.get_es()

        # TODO: change this to a bulk index
        for item in data:
            es.index(index, doctype, item, id=item['id'])

        cls.refresh()

    @classmethod
    def cleanup_index(cls):
        es = cls.get_es()
        try:
            es.indices.delete(cls.index_name)
        except elasticsearch.NotFoundError:
            pass

    @classmethod
    def refresh(cls):
        """Refresh index after indexing.

        This refreshes the index specified by `self.index_name`.

        """
        cls.get_es().indices.refresh(cls.index_name)
        cls.get_es().cluster.health(wait_for_status='yellow')

def facet_counts_dict(qs, field):
    return dict((t['term'], t['count']) for t in qs.facet_counts()[field])

def require_version(minimum_version):
    """Skip the test if the Elasticsearch version is less than specified.

    :arg minimum_version: string; the minimum Elasticsearch version required

    """

    def decorated(test):
        """Decorator to only run the test if ES version is greater or
        equal than specified.

        """

        @wraps(test)
        def test_with_version(self):
            "Only run the test if ES version is not less than specified."
            actual_version = self.get_es().info()['version']['number']

            if LooseVersion(actual_version) >= LooseVersion(minimum_version):
                test(self)
            else:
                raise SkipTest

        return test_with_version

    return decorated
<|endoftext|>