<commit_before># coding=utf-8

from __future__ import unicode_literals

import re
import unittest

from decimal import Decimal
from ukpostcodeparser.parser import parse_uk_postcode

from faker import Faker
from faker.providers.address.de_DE import Provider as DeProvider
from faker.providers.address.el_GR import Provider as GrProvider
from faker.providers.address.en_AU import Provider as EnAuProvider
from faker.providers.address.en_CA import Provider as EnCaProvider
from faker.providers.address.en_US import Provider as EnUsProvider
from faker.providers.address.ja_JP import Provider as JaProvider
from faker.providers.address.ne_NP import Provider as NeProvider
from six import string_types

class TestDeDE(unittest.TestCase):
    """ Tests in addresses in the de_DE locale """

    def setUp(self):
        self.factory = Faker('de_DE')

    def test_city(self):
        city = self.factory.city()
        assert isinstance(city, string_types)
        assert city in DeProvider.cities

    def test_state(self):
        state = self.factory.state()
        assert isinstance(state, string_types)
        assert state in DeProvider.states

    def test_street_suffix_short(self):
        street_suffix_short = self.factory.street_suffix_short()
        assert isinstance(street_suffix_short, string_types)
        assert street_suffix_short in DeProvider.street_suffixes_short

    def test_street_suffix_long(self):
        street_suffix_long = self.factory.street_suffix_long()
        assert isinstance(street_suffix_long, string_types)
        assert street_suffix_long in DeProvider.street_suffixes_long

    def test_country(self):
        country = self.factory.country()
        assert isinstance(country, string_types)
        assert country in DeProvider.countries

class TestElGR(unittest.TestCase):
    """ Tests addresses in the el_GR locale """

    def setUp(self):
        self.factory = Faker('el_GR')

    def test_line_address(self):
        address = self.factory.line_address()
        assert isinstance(address, string_types)

    def test_city(self):
        city = self.factory.city()
        assert isinstance(city, string_types)
        assert city in GrProvider.cities

    def test_region(self):
        region = self.factory.region()
        assert isinstance(region, string_types)
        assert region in GrProvider.regions

    def test_latlng(self):
        latlng = self.factory.latlng()
        latitude = self.factory.latitude()
        longitude = self.factory.longitude()
        assert isinstance(latlng, tuple)
        assert isinstance(latitude, Decimal)
        assert isinstance(longitude, Decimal)

class TestEnAU(unittest.TestCase):
    """ Tests addresses in the en_AU locale """

    def setUp(self):
        self.factory = Faker('en_AU')

    def test_postcode(self):
        for _ in range(100):
            postcode = self.factory.postcode()
            assert re.match("\d{4}", postcode)

    def test_state(self):
        state = self.factory.state()
        assert isinstance(state, string_types)
        assert state in EnAuProvider.states

    def test_city_prefix(self):
        city_prefix = self.factory.city_prefix()
        assert isinstance(city_prefix, string_types)
        assert city_prefix in EnAuProvider.city_prefixes

    def test_state_abbr(self):
        state_abbr = self.factory.state_abbr()
        assert isinstance(state_abbr, string_types)
        assert state_abbr in EnAuProvider.states_abbr
        self.assertTrue(state_abbr.isupper())

class TestEnCA(unittest.TestCase):
    """ Tests addresses in en_CA locale """

    def setUp(self):
        self.factory = Faker('en_CA')

    def test_postalcode(self):
        for _ in range(100):
            postalcode = self.factory.postalcode()
            assert re.match("[A-Z][0-9][A-Z] ?[0-9][A-Z][0-9]",
                            postalcode)

    def test_postal_code_letter(self):
        postal_code_letter = self.factory.postal_code_letter()
        assert re.match("[A-Z]", postal_code_letter)

    def test_province(self):
        province = self.factory.province()
        assert isinstance(province, string_types)
        assert province in EnCaProvider.provinces

    def test_province_abbr(self):
        province_abbr = self.factory.province_abbr()
        assert isinstance(province_abbr, string_types)
        assert province_abbr in EnCaProvider.provinces_abbr

    def test_city_prefix(self):
        city_prefix = self.factory.city_prefix()
        assert isinstance(city_prefix, string_types)
        assert city_prefix in EnCaProvider.city_prefixes

    def test_secondary_address(self):
        secondary_address = self.factory.secondary_address()
        assert isinstance(secondary_address, string_types)

class TestEnGB(unittest.TestCase):
    """ Tests addresses in the en_GB locale """

    def setUp(self):
        self.factory = Faker('en_GB')

    def test_postcode(self):
        for _ in range(100):
            assert isinstance(parse_uk_postcode(self.factory.postcode()), tuple)

class TestEnUS(unittest.TestCase):
    """ Tests addresses in the en_US locale """

    def setUp(self):
        self.factory = Faker('en_US')

    def test_city_prefix(self):
        city_prefix = self.factory.city_prefix()
        assert isinstance(city_prefix, string_types)
        assert city_prefix in EnUsProvider.city_prefixes

    def test_state(self):
        state = self.factory.state()
        assert isinstance(state, string_types)
        assert state in EnUsProvider.states

    def test_state_abbr(self):
        state_abbr = self.factory.state_abbr()
        assert isinstance(state_abbr, string_types)
        assert state_abbr in EnUsProvider.states_abbr

    def test_zipcode(self):
        for _ in range(100):
            zipcode = self.factory.zipcode()
            assert re.match("\d{5}", zipcode)

class TestHuHU(unittest.TestCase):
    """ Tests addresses in the hu_HU locale """

    def setUp(self):
        self.factory = Faker('hu_HU')

    def test_postcode_first_digit(self):
        # Hungarian postcodes begin with 'H-' followed by 4 digits.
        # The first digit may not begin with a zero.
        for _ in range(100):
            pcd = self.factory.postcode()
            assert pcd[2] > "0"

    def test_street_address(self):
        """ Tests the street address in the hu_HU locale """
        address = self.factory.address()
        assert isinstance(address, string_types)
        address_with_county = self.factory.street_address_with_county()
        assert isinstance(address_with_county, string_types)

class TestJaJP(unittest.TestCase):
    """ Tests addresses in the ja_JP locale """

    def setUp(self):
        self.factory = Faker('ja')

    def test_address(self):
        """ Test"""
        country = self.factory.country()
        assert isinstance(country, string_types)
        assert country in JaProvider.countries

        prefecture = self.factory.prefecture()
        assert isinstance(prefecture, string_types)
        assert prefecture in JaProvider.prefectures

        city = self.factory.city()
        assert isinstance(city, string_types)
        assert city in JaProvider.cities

        town = self.factory.town()
        assert isinstance(town, string_types)
        assert town in JaProvider.towns

        chome = self.factory.chome()
        assert isinstance(chome, string_types)
        assert re.match("\d{1,2}丁目", chome)

        ban = self.factory.ban()
        assert isinstance(ban, string_types)
        assert re.match("\d{1,2}番", ban)

        gou = self.factory.gou()
        assert isinstance(gou, string_types)
        assert re.match("\d{1,2}号", gou)

        building_name = self.factory.building_name()
        assert isinstance(building_name, string_types)
        assert building_name in JaProvider.building_names

        zipcode = self.factory.zipcode()
        assert isinstance(zipcode, string_types)
        assert re.match("\d{3}-\d{4}", zipcode)

        address = self.factory.address()
        assert isinstance(address, string_types)

class TestNeNP(unittest.TestCase):
    """ Tests addresses in the ne_NP locale """

    def setUp(self):
        self.factory = Faker('ne_NP')

    def test_address(self):
        """ Tests the street address in ne_NP locale """
        country = self.factory.country()
        assert isinstance(country, string_types)
        assert country in NeProvider.countries

        district = self.factory.district()
        assert isinstance(district, string_types)
        assert district in NeProvider.districts

        city = self.factory.city()
        assert isinstance(city, string_types)
        assert city in NeProvider.cities

class TestNoNO(unittest.TestCase):
    """ Tests the street address in no_NO locale """

    def setUp(self):
        self.factory = Faker('no_NO')

    def test_postcode(self):
        for _ in range(100):
            self.assertTrue(re.match(r'^[0-9]{4}$', self.factory.postcode()))

    def test_city_suffix(self):
        suffix = self.factory.city_suffix()
        assert isinstance(suffix, string_types)

    def test_street_suffix(self):
        suffix = self.factory.street_suffix()
        assert isinstance(suffix, string_types)

    def test_address(self):
        address = self.factory.address()
        assert isinstance(address, string_types)

class TestZhTW(unittest.TestCase):
    """ Tests addresses in the zh_tw locale """

    def setUp(self):
        self.factory = Faker('zh_TW')

    def test_address(self):
        country = self.factory.country()
        assert isinstance(country, string_types)

        street = self.factory.street_name()
        assert isinstance(street, string_types)

        city = self.factory.city()
        assert isinstance(city, string_types)

        address = self.factory.address()
        assert isinstance(address, string_types)

class TestZhCN(unittest.TestCase):
    """ Tests addresses in the zh_cn locale """

    def setUp(self):
        self.factory = Faker('zh_CN')

    def test_address(self):
        country = self.factory.country()
        assert isinstance(country, string_types)

        street = self.factory.street_name()
        assert isinstance(street, string_types)

        city = self.factory.street_address()
        assert isinstance(city, string_types)

        province = self.factory.province()
        assert isinstance(province, string_types)

        district = self.factory.district()
        assert isinstance(district, string_types)

        address = self.factory.address()
        assert isinstance(address, string_types)

        for _ in range(100):
            self.assertTrue(re.match(r'\d{5}', self.factory.postcode()))
<commit_msg>Introduce military ship test<commit_after># coding=utf-8

from __future__ import unicode_literals

import re
import unittest

from decimal import Decimal
from ukpostcodeparser.parser import parse_uk_postcode

from faker import Faker
from faker.providers.address.de_DE import Provider as DeProvider
from faker.providers.address.el_GR import Provider as GrProvider
from faker.providers.address.en_AU import Provider as EnAuProvider
from faker.providers.address.en_CA import Provider as EnCaProvider
from faker.providers.address.en_US import Provider as EnUsProvider
from faker.providers.address.ja_JP import Provider as JaProvider
from faker.providers.address.ne_NP import Provider as NeProvider
from six import string_types

class TestDeDE(unittest.TestCase):
    """ Tests in addresses in the de_DE locale """

    def setUp(self):
        self.factory = Faker('de_DE')

    def test_city(self):
        city = self.factory.city()
        assert isinstance(city, string_types)
        assert city in DeProvider.cities

    def test_state(self):
        state = self.factory.state()
        assert isinstance(state, string_types)
        assert state in DeProvider.states

    def test_street_suffix_short(self):
        street_suffix_short = self.factory.street_suffix_short()
        assert isinstance(street_suffix_short, string_types)
        assert street_suffix_short in DeProvider.street_suffixes_short

    def test_street_suffix_long(self):
        street_suffix_long = self.factory.street_suffix_long()
        assert isinstance(street_suffix_long, string_types)
        assert street_suffix_long in DeProvider.street_suffixes_long

    def test_country(self):
        country = self.factory.country()
        assert isinstance(country, string_types)
        assert country in DeProvider.countries

class TestElGR(unittest.TestCase):
    """ Tests addresses in the el_GR locale """

    def setUp(self):
        self.factory = Faker('el_GR')

    def test_line_address(self):
        address = self.factory.line_address()
        assert isinstance(address, string_types)

    def test_city(self):
        city = self.factory.city()
        assert isinstance(city, string_types)
        assert city in GrProvider.cities

    def test_region(self):
        region = self.factory.region()
        assert isinstance(region, string_types)
        assert region in GrProvider.regions

    def test_latlng(self):
        latlng = self.factory.latlng()
        latitude = self.factory.latitude()
        longitude = self.factory.longitude()
        assert isinstance(latlng, tuple)
        assert isinstance(latitude, Decimal)
        assert isinstance(longitude, Decimal)

class TestEnAU(unittest.TestCase):
    """ Tests addresses in the en_AU locale """

    def setUp(self):
        self.factory = Faker('en_AU')

    def test_postcode(self):
        for _ in range(100):
            postcode = self.factory.postcode()
            assert re.match("\d{4}", postcode)

    def test_state(self):
        state = self.factory.state()
        assert isinstance(state, string_types)
        assert state in EnAuProvider.states

    def test_city_prefix(self):
        city_prefix = self.factory.city_prefix()
        assert isinstance(city_prefix, string_types)
        assert city_prefix in EnAuProvider.city_prefixes

    def test_state_abbr(self):
        state_abbr = self.factory.state_abbr()
        assert isinstance(state_abbr, string_types)
        assert state_abbr in EnAuProvider.states_abbr
        self.assertTrue(state_abbr.isupper())

class TestEnCA(unittest.TestCase):
    """ Tests addresses in en_CA locale """

    def setUp(self):
        self.factory = Faker('en_CA')

    def test_postalcode(self):
        for _ in range(100):
            postalcode = self.factory.postalcode()
            assert re.match("[A-Z][0-9][A-Z] ?[0-9][A-Z][0-9]",
                            postalcode)

    def test_postal_code_letter(self):
        postal_code_letter = self.factory.postal_code_letter()
        assert re.match("[A-Z]", postal_code_letter)

    def test_province(self):
        province = self.factory.province()
        assert isinstance(province, string_types)
        assert province in EnCaProvider.provinces

    def test_province_abbr(self):
        province_abbr = self.factory.province_abbr()
        assert isinstance(province_abbr, string_types)
        assert province_abbr in EnCaProvider.provinces_abbr

    def test_city_prefix(self):
        city_prefix = self.factory.city_prefix()
        assert isinstance(city_prefix, string_types)
        assert city_prefix in EnCaProvider.city_prefixes

    def test_secondary_address(self):
        secondary_address = self.factory.secondary_address()
        assert isinstance(secondary_address, string_types)

class TestEnGB(unittest.TestCase):
    """ Tests addresses in the en_GB locale """

    def setUp(self):
        self.factory = Faker('en_GB')

    def test_postcode(self):
        for _ in range(100):
            assert isinstance(parse_uk_postcode(self.factory.postcode()), tuple)

class TestEnUS(unittest.TestCase):
    """ Tests addresses in the en_US locale """

    def setUp(self):
        self.factory = Faker('en_US')

    def test_city_prefix(self):
        city_prefix = self.factory.city_prefix()
        assert isinstance(city_prefix, string_types)
        assert city_prefix in EnUsProvider.city_prefixes

    def test_state(self):
        state = self.factory.state()
        assert isinstance(state, string_types)
        assert state in EnUsProvider.states

    def test_state_abbr(self):
        state_abbr = self.factory.state_abbr()
        assert isinstance(state_abbr, string_types)
        assert state_abbr in EnUsProvider.states_abbr

    def test_zipcode(self):
        for _ in range(100):
            zipcode = self.factory.zipcode()
            assert re.match("\d{5}", zipcode)

    def test_zipcode_plus4(self):
        for _ in range(100):
            zipcode_plus4 = self.factory.zipcode_plus4()
            assert re.match("\d{5}(-\d{4})", zipcode_plus4)

    def test_military_ship(self):
        military_ship = self.factory.military_ship()
        assert isinstance(military_ship, string_types)
        assert military_ship in EnUsProvider.military_ship_prefix
        assert re.match("[A-Z]", military_ship)

class TestHuHU(unittest.TestCase):
    """ Tests addresses in the hu_HU locale """

    def setUp(self):
        self.factory = Faker('hu_HU')

    def test_postcode_first_digit(self):
        # Hungarian postcodes begin with 'H-' followed by 4 digits.
        # The first digit may not begin with a zero.
        for _ in range(100):
            pcd = self.factory.postcode()
            assert pcd[2] > "0"

    def test_street_address(self):
        """ Tests the street address in the hu_HU locale """
        address = self.factory.address()
        assert isinstance(address, string_types)
        address_with_county = self.factory.street_address_with_county()
        assert isinstance(address_with_county, string_types)

class TestJaJP(unittest.TestCase):
    """ Tests addresses in the ja_JP locale """

    def setUp(self):
        self.factory = Faker('ja')

    def test_address(self):
        """ Test"""
        country = self.factory.country()
        assert isinstance(country, string_types)
        assert country in JaProvider.countries

        prefecture = self.factory.prefecture()
        assert isinstance(prefecture, string_types)
        assert prefecture in JaProvider.prefectures

        city = self.factory.city()
        assert isinstance(city, string_types)
        assert city in JaProvider.cities

        town = self.factory.town()
        assert isinstance(town, string_types)
        assert town in JaProvider.towns

        chome = self.factory.chome()
        assert isinstance(chome, string_types)
        assert re.match("\d{1,2}丁目", chome)

        ban = self.factory.ban()
        assert isinstance(ban, string_types)
        assert re.match("\d{1,2}番", ban)

        gou = self.factory.gou()
        assert isinstance(gou, string_types)
        assert re.match("\d{1,2}号", gou)

        building_name = self.factory.building_name()
        assert isinstance(building_name, string_types)
        assert building_name in JaProvider.building_names

        zipcode = self.factory.zipcode()
        assert isinstance(zipcode, string_types)
        assert re.match("\d{3}-\d{4}", zipcode)

        address = self.factory.address()
        assert isinstance(address, string_types)

class TestNeNP(unittest.TestCase):
    """ Tests addresses in the ne_NP locale """

    def setUp(self):
        self.factory = Faker('ne_NP')

    def test_address(self):
        """ Tests the street address in ne_NP locale """
        country = self.factory.country()
        assert isinstance(country, string_types)
        assert country in NeProvider.countries

        district = self.factory.district()
        assert isinstance(district, string_types)
        assert district in NeProvider.districts

        city = self.factory.city()
        assert isinstance(city, string_types)
        assert city in NeProvider.cities

class TestNoNO(unittest.TestCase):
    """ Tests the street address in no_NO locale """

    def setUp(self):
        self.factory = Faker('no_NO')

    def test_postcode(self):
        for _ in range(100):
            self.assertTrue(re.match(r'^[0-9]{4}$', self.factory.postcode()))

    def test_city_suffix(self):
        suffix = self.factory.city_suffix()
        assert isinstance(suffix, string_types)

    def test_street_suffix(self):
        suffix = self.factory.street_suffix()
        assert isinstance(suffix, string_types)

    def test_address(self):
        address = self.factory.address()
        assert isinstance(address, string_types)

class TestZhTW(unittest.TestCase):
    """ Tests addresses in the zh_tw locale """

    def setUp(self):
        self.factory = Faker('zh_TW')

    def test_address(self):
        country = self.factory.country()
        assert isinstance(country, string_types)

        street = self.factory.street_name()
        assert isinstance(street, string_types)

        city = self.factory.city()
        assert isinstance(city, string_types)

        address = self.factory.address()
        assert isinstance(address, string_types)

class TestZhCN(unittest.TestCase):
    """ Tests addresses in the zh_cn locale """

    def setUp(self):
        self.factory = Faker('zh_CN')

    def test_address(self):
        country = self.factory.country()
        assert isinstance(country, string_types)

        street = self.factory.street_name()
        assert isinstance(street, string_types)

        city = self.factory.street_address()
        assert isinstance(city, string_types)

        province = self.factory.province()
        assert isinstance(province, string_types)

        district = self.factory.district()
        assert isinstance(district, string_types)

        address = self.factory.address()
        assert isinstance(address, string_types)

        for _ in range(100):
            self.assertTrue(re.match(r'\d{5}', self.factory.postcode()))
<|endoftext|>


<commit_before>import asyncio
import logging
from contextlib import suppress
from functools import partial
from typing import Callable, Type

import aiomisc
import aiormq.exceptions
import pytest
import shortuuid
from aiomisc_pytest.pytest_plugin import TCPProxy
from pamqp.exceptions import AMQPFrameError
from yarl import URL

import aio_pika
from aio_pika.exceptions import QueueEmpty
from aio_pika.message import IncomingMessage, Message
from aio_pika.robust_channel import RobustChannel
from aio_pika.robust_connection import RobustConnection
from aio_pika.robust_queue import RobustQueue
from tests import get_random_name

@pytest.fixture
async def proxy(tcp_proxy: Type[TCPProxy], amqp_direct_url: URL):
    p = tcp_proxy(amqp_direct_url.host, amqp_direct_url.port)

    await p.start()
    try:
        yield p
    finally:
        await p.close()

@pytest.fixture
def amqp_url(amqp_direct_url, proxy: TCPProxy):
    return amqp_direct_url.with_host(
        proxy.proxy_host,
    ).with_port(
        proxy.proxy_port,
    ).update_query(
        reconnect_interval=1,
        heartbeat=1,
    )

@pytest.fixture
def proxy_port(aiomisc_unused_port_factory) -> int:
    return aiomisc_unused_port_factory()

@pytest.fixture(scope="module")
def connection_fabric():
    return aio_pika.connect_robust

@pytest.fixture
def create_direct_connection(loop, amqp_direct_url):
    return partial(
        aio_pika.connect,
        amqp_direct_url.update_query(
            name=amqp_direct_url.query["name"] + "::direct",
            heartbeat=30,
        ),
        loop=loop,
    )

@pytest.fixture
def create_connection(connection_fabric, loop, amqp_url):
    return partial(connection_fabric, amqp_url, loop=loop)

@pytest.fixture
async def direct_connection(create_direct_connection) -> aio_pika.Connection:
    async with await create_direct_connection() as conn:
        yield conn

async def test_channel_fixture(channel: aio_pika.RobustChannel):
    assert isinstance(channel, aio_pika.RobustChannel)

async def test_connection_fixture(connection: aio_pika.RobustConnection):
    assert isinstance(connection, aio_pika.RobustConnection)

def test_amqp_url_is_not_direct(amqp_url, amqp_direct_url):
    assert amqp_url != amqp_direct_url

async def test_set_qos(channel: aio_pika.Channel):
    await channel.set_qos(prefetch_count=1)

async def test_revive_passive_queue_on_reconnect(
    create_connection, direct_connection, proxy: TCPProxy,
):
    client = await create_connection()
    assert isinstance(client, RobustConnection)

    reconnect_event = asyncio.Event()
    reconnect_count = 0

    def reconnect_callback(sender, conn):
        nonlocal reconnect_count
        reconnect_count += 1
        reconnect_event.set()
        reconnect_event.clear()

    client.reconnect_callbacks.add(reconnect_callback)

    queue_name = get_random_name()
    channel = await client.channel()
    assert isinstance(channel, RobustChannel)

    direct_channel = await direct_connection.channel()

    direct_queue = await direct_channel.declare_queue(
        queue_name, auto_delete=True, passive=False,
    )

    queue2 = await channel.declare_queue(
        direct_queue.name, passive=True, auto_delete=False,
    )
    assert isinstance(queue2, RobustQueue)

    await proxy.disconnect_all()
    await reconnect_event.wait()

    assert reconnect_count == 1

    with suppress(asyncio.TimeoutError):
        await asyncio.wait_for(
            reconnect_event.wait(), client.reconnect_interval * 2,
        )

    assert reconnect_count == 1

@aiomisc.timeout(30)
async def test_robust_reconnect(
    create_connection, direct_connection,
    proxy: TCPProxy, loop, add_cleanup: Callable,
):
    read_conn = await create_connection()   # type: aio_pika.RobustConnection

    reconnect_event = asyncio.Event()
    read_conn.reconnect_callbacks.add(
        lambda *_: reconnect_event.set(),
    )

    assert isinstance(read_conn, aio_pika.RobustConnection)

    async with read_conn, direct_connection:
        read_channel = await read_conn.channel()
        write_channel = await direct_connection.channel()

        assert isinstance(read_channel, aio_pika.RobustChannel)

        qname = get_random_name("robust", "proxy", "shared")

        async with read_channel, write_channel:
            shared = []

            # Declaring temporary queue
            queue = await write_channel.declare_queue(
                qname,
                auto_delete=False,
                durable=True,
            )

            consumer_event = asyncio.Event()

            async def reader(queue_name):
                nonlocal shared

                try:
                    queue = await read_channel.declare_queue(
                        name=queue_name, passive=True,
                    )

                    async with queue.iterator() as q:
                        loop.call_soon(consumer_event.set)

                        async for message in q:
                            shared.append(message)
                            await message.ack()
                finally:
                    logging.info("Exit reader task")

            try:
                reader_task = loop.create_task(reader(queue.name))

                await consumer_event.wait()
                logging.info("Disconnect all clients")
                with proxy.slowdown(1, 1):
                    for i in range(5):
                        await write_channel.default_exchange.publish(
                            Message(str(i).encode()), queue.name,
                        )

                    await proxy.disconnect_all()

                    # noinspection PyTypeChecker
                    with pytest.raises(AMQPFrameError):
                        await read_conn.channel()

                logging.info("Waiting reconnect")
                await reconnect_event.wait()

                logging.info("Waiting connections")
                await asyncio.wait_for(read_conn.ready(), timeout=20)

                for i in range(5, 10):
                    await write_channel.default_exchange.publish(
                        Message(str(i).encode()), queue.name,
                    )

                while len(shared) < 10:
                    await asyncio.sleep(0.1)

                assert len(shared) == 10

                reader_task.cancel()
                await asyncio.gather(reader_task, return_exceptions=True)

                with pytest.raises(QueueEmpty):
                    await queue.get(timeout=0.5)
            finally:
                await queue.purge()
                await queue.delete()

async def test_channel_locked_resource2(connection: aio_pika.RobustConnection):
    ch1 = await connection.channel()
    ch2 = await connection.channel()

    qname = get_random_name("channel", "locked", "resource")

    q1 = await ch1.declare_queue(qname, exclusive=True, robust=False)
    await q1.consume(print, exclusive=True)

    with pytest.raises(aiormq.exceptions.ChannelAccessRefused):
        q2 = await ch2.declare_queue(qname, exclusive=True, robust=False)
        await q2.consume(print, exclusive=True)

async def test_channel_close_when_exclusive_queue(
    create_connection, create_direct_connection, proxy: TCPProxy, loop,
):
    logging.info("Creating connections")
    direct_conn, proxy_conn = await asyncio.gather(
        create_direct_connection(), create_connection(),
    )

    logging.info("Creating channels")
    direct_channel, proxy_channel = await asyncio.gather(
        direct_conn.channel(), proxy_conn.channel(),
    )

    reconnect_event = asyncio.Event()
    proxy_conn.reconnect_callbacks.add(
        lambda *_: reconnect_event.set(), weak=False,
    )

    qname = get_random_name("robust", "exclusive", "queue")

    logging.info("Declaring exclusing queue: %s", qname)
    proxy_queue = await proxy_channel.declare_queue(
        qname, exclusive=True, durable=True,
    )

    logging.info("Disconnecting all proxy connections")
    await proxy.disconnect_all()
    await asyncio.sleep(0.5)

    logging.info("Declaring exclusive queue through direct channel")
    await direct_channel.declare_queue(
        qname, exclusive=True, durable=True,
    )

    async def close_after(delay, closer):
        await asyncio.sleep(delay)
        logging.info("Disconnecting direct connection")
        await closer()
        logging.info("Closed")

    await loop.create_task(close_after(5, direct_conn.close))

    # reconnect fired
    await reconnect_event.wait()

    # Wait method ready
    await proxy_conn.connected.wait()
    await proxy_queue.delete()

async def test_context_process_abrupt_channel_close(
    connection: aio_pika.RobustConnection,
    declare_exchange: Callable,
    declare_queue: Callable,
):
    # https://github.com/mosquito/aio-pika/issues/302
    queue_name = get_random_name("test_connection")
    routing_key = get_random_name("rounting_key")

    channel = await connection.channel()
    exchange = await declare_exchange(
        "direct", auto_delete=True, channel=channel,
    )
    queue = await declare_queue(queue_name, auto_delete=True, channel=channel)

    await queue.bind(exchange, routing_key)
    body = bytes(shortuuid.uuid(), "utf-8")

    await exchange.publish(
        Message(body, content_type="text/plain", headers={"foo": "bar"}),
        routing_key,
    )

    incoming_message = await queue.get(timeout=5)
    # close aiormq channel to emulate abrupt connection/channel close
    await channel.channel.close()
    with pytest.raises(aiormq.exceptions.ChannelInvalidStateError):
        async with incoming_message.process():
            # emulate some activity on closed channel
            await channel.channel.basic_publish(
                "dummy", exchange="", routing_key="non_existent",
            )

    # emulate connection/channel restoration of connect_robust
    await channel.reopen()

    # cleanup queue
    incoming_message = await queue.get(timeout=5)
    async with incoming_message.process():
        pass
    await queue.unbind(exchange, routing_key)

@aiomisc.timeout(10)
async def test_robust_duplicate_queue(
    connection: aio_pika.RobustConnection,
    direct_connection: aio_pika.Connection,
    declare_exchange: Callable,
    declare_queue: Callable,
    proxy: TCPProxy,
    create_task: Callable,
):
    queue_name = get_random_name("test")

    channel = await connection.channel()
    direct_channel = await direct_connection.channel()

    reconnect_event = asyncio.Event()
    shared_condition = asyncio.Condition()

    connection.reconnect_callbacks.add(
        lambda *_: reconnect_event.set(),
    )

    shared = {}

    # noinspection PyShadowingNames
    async def reader(queue: aio_pika.Queue):
        nonlocal shared

        async with queue.iterator() as q:
            async for message in q:
                message: IncomingMessage
                # https://www.rabbitmq.com/confirms.html#automatic-requeueing
                async with shared_condition:
                    shared[message.message_id] = message
                    shared_condition.notify_all()
                    await message.ack()

    queue = await declare_queue(
        queue_name, channel=channel, cleanup=False,
    )

    create_task(reader(queue))

    for x in range(5):
        await direct_channel.default_exchange.publish(
            aio_pika.Message(b"1234567890", message_id=f"0-{x}"), queue_name,
        )

    async with shared_condition:
        await asyncio.wait_for(
            shared_condition.wait_for(lambda: len(shared) == 5),
            timeout=5,
        )

    logging.info("Disconnect all clients")
    await proxy.disconnect_all()

    assert len(shared) == 5, shared

    for x in range(5):
        await direct_channel.default_exchange.publish(
            Message(b"1234567890", message_id=f"1-{x}"), queue_name,
        )

    await asyncio.wait_for(reconnect_event.wait(), timeout=5)

    logging.info("Waiting connections")
    await channel.connection.ready()

    async with shared_condition:
        await asyncio.wait_for(
            shared_condition.wait_for(lambda: len(shared) == 10),
            timeout=5,
        )

    assert len(shared) == 10

@aiomisc.timeout(10)
async def test_channel_reconnect(
    connection_fabric, loop, amqp_url, proxy: TCPProxy, add_cleanup: Callable,
):
    heartbeat = 2
    amqp_url = amqp_url.update_query(heartbeat=heartbeat)

    on_reconnect = asyncio.Event()

    conn = await connection_fabric(amqp_url, loop=loop)
    assert isinstance(conn, aio_pika.RobustConnection)

    conn.reconnect_callbacks.add(lambda *_: on_reconnect.set(), weak=False)

    async with conn:
        channel = await conn.channel()
        assert isinstance(channel, aio_pika.RobustChannel)

        async with channel:
            await channel.set_qos(0)
            await channel.set_qos(1)

            with pytest.raises(asyncio.TimeoutError):
                with proxy.slowdown(1, 1):
                    await channel.set_qos(0, timeout=0.5)

            await on_reconnect.wait()
            await channel.set_qos(0)
            await channel.set_qos(1)
<commit_msg>add failing test<commit_after>import asyncio
import logging
from contextlib import suppress
from functools import partial
from typing import Callable, Type

import aiomisc
import aiormq.exceptions
import pytest
import shortuuid
from aiomisc_pytest.pytest_plugin import TCPProxy
from pamqp.exceptions import AMQPFrameError
from yarl import URL

import aio_pika
from aio_pika.exceptions import QueueEmpty
from aio_pika.message import IncomingMessage, Message
from aio_pika.robust_channel import RobustChannel
from aio_pika.robust_connection import RobustConnection
from aio_pika.robust_queue import RobustQueue
from tests import get_random_name

@pytest.fixture
async def proxy(tcp_proxy: Type[TCPProxy], amqp_direct_url: URL):
    p = tcp_proxy(amqp_direct_url.host, amqp_direct_url.port)

    await p.start()
    try:
        yield p
    finally:
        await p.close()

@pytest.fixture
def amqp_url(amqp_direct_url, proxy: TCPProxy):
    return amqp_direct_url.with_host(
        proxy.proxy_host,
    ).with_port(
        proxy.proxy_port,
    ).update_query(
        reconnect_interval=1,
        heartbeat=1,
    )

@pytest.fixture
def proxy_port(aiomisc_unused_port_factory) -> int:
    return aiomisc_unused_port_factory()

@pytest.fixture(scope="module")
def connection_fabric():
    return aio_pika.connect_robust

@pytest.fixture
def create_direct_connection(loop, amqp_direct_url):
    return partial(
        aio_pika.connect,
        amqp_direct_url.update_query(
            name=amqp_direct_url.query["name"] + "::direct",
            heartbeat=30,
        ),
        loop=loop,
    )

@pytest.fixture
def create_connection(connection_fabric, loop, amqp_url):
    return partial(connection_fabric, amqp_url, loop=loop)

@pytest.fixture
async def direct_connection(create_direct_connection) -> aio_pika.Connection:
    async with await create_direct_connection() as conn:
        yield conn

async def test_channel_fixture(channel: aio_pika.RobustChannel):
    assert isinstance(channel, aio_pika.RobustChannel)

async def test_connection_fixture(connection: aio_pika.RobustConnection):
    assert isinstance(connection, aio_pika.RobustConnection)

def test_amqp_url_is_not_direct(amqp_url, amqp_direct_url):
    assert amqp_url != amqp_direct_url

async def test_set_qos(channel: aio_pika.Channel):
    await channel.set_qos(prefetch_count=1)

async def test_revive_passive_queue_on_reconnect(
    create_connection, direct_connection, proxy: TCPProxy,
):
    client = await create_connection()
    assert isinstance(client, RobustConnection)

    reconnect_event = asyncio.Event()
    reconnect_count = 0

    def reconnect_callback(sender, conn):
        nonlocal reconnect_count
        reconnect_count += 1
        reconnect_event.set()
        reconnect_event.clear()

    client.reconnect_callbacks.add(reconnect_callback)

    queue_name = get_random_name()
    channel = await client.channel()
    assert isinstance(channel, RobustChannel)

    direct_channel = await direct_connection.channel()

    direct_queue = await direct_channel.declare_queue(
        queue_name, auto_delete=True, passive=False,
    )

    queue2 = await channel.declare_queue(
        direct_queue.name, passive=True, auto_delete=False,
    )
    assert isinstance(queue2, RobustQueue)

    await proxy.disconnect_all()
    await reconnect_event.wait()

    assert reconnect_count == 1

    with suppress(asyncio.TimeoutError):
        await asyncio.wait_for(
            reconnect_event.wait(), client.reconnect_interval * 2,
        )

    assert reconnect_count == 1

@aiomisc.timeout(30)
async def test_robust_reconnect(
    create_connection, direct_connection,
    proxy: TCPProxy, loop, add_cleanup: Callable,
):
    read_conn = await create_connection()   # type: aio_pika.RobustConnection

    reconnect_event = asyncio.Event()
    read_conn.reconnect_callbacks.add(
        lambda *_: reconnect_event.set(),
    )

    assert isinstance(read_conn, aio_pika.RobustConnection)

    async with read_conn, direct_connection:
        read_channel = await read_conn.channel()
        write_channel = await direct_connection.channel()

        assert isinstance(read_channel, aio_pika.RobustChannel)

        qname = get_random_name("robust", "proxy", "shared")

        async with read_channel, write_channel:
            shared = []

            # Declaring temporary queue
            queue = await write_channel.declare_queue(
                qname,
                auto_delete=False,
                durable=True,
            )

            consumer_event = asyncio.Event()

            async def reader(queue_name):
                nonlocal shared

                try:
                    queue = await read_channel.declare_queue(
                        name=queue_name, passive=True,
                    )

                    async with queue.iterator() as q:
                        loop.call_soon(consumer_event.set)

                        async for message in q:
                            shared.append(message)
                            await message.ack()
                finally:
                    logging.info("Exit reader task")

            try:
                reader_task = loop.create_task(reader(queue.name))

                await consumer_event.wait()
                logging.info("Disconnect all clients")
                with proxy.slowdown(1, 1):
                    for i in range(5):
                        await write_channel.default_exchange.publish(
                            Message(str(i).encode()), queue.name,
                        )

                    await proxy.disconnect_all()

                    # noinspection PyTypeChecker
                    with pytest.raises(AMQPFrameError):
                        await read_conn.channel()

                logging.info("Waiting reconnect")
                await reconnect_event.wait()

                logging.info("Waiting connections")
                await asyncio.wait_for(read_conn.ready(), timeout=20)

                for i in range(5, 10):
                    await write_channel.default_exchange.publish(
                        Message(str(i).encode()), queue.name,
                    )

                while len(shared) < 10:
                    await asyncio.sleep(0.1)

                assert len(shared) == 10

                reader_task.cancel()
                await asyncio.gather(reader_task, return_exceptions=True)

                with pytest.raises(QueueEmpty):
                    await queue.get(timeout=0.5)
            finally:
                await queue.purge()
                await queue.delete()

async def test_channel_locked_resource2(connection: aio_pika.RobustConnection):
    ch1 = await connection.channel()
    ch2 = await connection.channel()

    qname = get_random_name("channel", "locked", "resource")

    q1 = await ch1.declare_queue(qname, exclusive=True, robust=False)
    await q1.consume(print, exclusive=True)

    with pytest.raises(aiormq.exceptions.ChannelAccessRefused):
        q2 = await ch2.declare_queue(qname, exclusive=True, robust=False)
        await q2.consume(print, exclusive=True)

async def test_channel_close_when_exclusive_queue(
    create_connection, create_direct_connection, proxy: TCPProxy, loop,
):
    logging.info("Creating connections")
    direct_conn, proxy_conn = await asyncio.gather(
        create_direct_connection(), create_connection(),
    )

    logging.info("Creating channels")
    direct_channel, proxy_channel = await asyncio.gather(
        direct_conn.channel(), proxy_conn.channel(),
    )

    reconnect_event = asyncio.Event()
    proxy_conn.reconnect_callbacks.add(
        lambda *_: reconnect_event.set(), weak=False,
    )

    qname = get_random_name("robust", "exclusive", "queue")

    logging.info("Declaring exclusing queue: %s", qname)
    proxy_queue = await proxy_channel.declare_queue(
        qname, exclusive=True, durable=True,
    )

    logging.info("Disconnecting all proxy connections")
    await proxy.disconnect_all()
    await asyncio.sleep(0.5)

    logging.info("Declaring exclusive queue through direct channel")
    await direct_channel.declare_queue(
        qname, exclusive=True, durable=True,
    )

    async def close_after(delay, closer):
        await asyncio.sleep(delay)
        logging.info("Disconnecting direct connection")
        await closer()
        logging.info("Closed")

    await loop.create_task(close_after(5, direct_conn.close))

    # reconnect fired
    await reconnect_event.wait()

    # Wait method ready
    await proxy_conn.connected.wait()
    await proxy_queue.delete()

async def test_context_process_abrupt_channel_close(
    connection: aio_pika.RobustConnection,
    declare_exchange: Callable,
    declare_queue: Callable,
):
    # https://github.com/mosquito/aio-pika/issues/302
    queue_name = get_random_name("test_connection")
    routing_key = get_random_name("rounting_key")

    channel = await connection.channel()
    exchange = await declare_exchange(
        "direct", auto_delete=True, channel=channel,
    )
    queue = await declare_queue(queue_name, auto_delete=True, channel=channel)

    await queue.bind(exchange, routing_key)
    body = bytes(shortuuid.uuid(), "utf-8")

    await exchange.publish(
        Message(body, content_type="text/plain", headers={"foo": "bar"}),
        routing_key,
    )

    incoming_message = await queue.get(timeout=5)
    # close aiormq channel to emulate abrupt connection/channel close
    await channel.channel.close()
    with pytest.raises(aiormq.exceptions.ChannelInvalidStateError):
        async with incoming_message.process():
            # emulate some activity on closed channel
            await channel.channel.basic_publish(
                "dummy", exchange="", routing_key="non_existent",
            )

    # emulate connection/channel restoration of connect_robust
    await channel.reopen()

    # cleanup queue
    incoming_message = await queue.get(timeout=5)
    async with incoming_message.process():
        pass
    await queue.unbind(exchange, routing_key)

@aiomisc.timeout(10)
async def test_robust_duplicate_queue(
    connection: aio_pika.RobustConnection,
    direct_connection: aio_pika.Connection,
    declare_exchange: Callable,
    declare_queue: Callable,
    proxy: TCPProxy,
    create_task: Callable,
):
    queue_name = get_random_name("test")

    channel = await connection.channel()
    direct_channel = await direct_connection.channel()

    reconnect_event = asyncio.Event()
    shared_condition = asyncio.Condition()

    connection.reconnect_callbacks.add(
        lambda *_: reconnect_event.set(),
    )

    shared = {}

    # noinspection PyShadowingNames
    async def reader(queue: aio_pika.Queue):
        nonlocal shared

        async with queue.iterator() as q:
            async for message in q:
                message: IncomingMessage
                # https://www.rabbitmq.com/confirms.html#automatic-requeueing
                async with shared_condition:
                    shared[message.message_id] = message
                    shared_condition.notify_all()
                    await message.ack()

    queue = await declare_queue(
        queue_name, channel=channel, cleanup=False,
    )

    create_task(reader(queue))

    for x in range(5):
        await direct_channel.default_exchange.publish(
            aio_pika.Message(b"1234567890", message_id=f"0-{x}"), queue_name,
        )

    async with shared_condition:
        await asyncio.wait_for(
            shared_condition.wait_for(lambda: len(shared) == 5),
            timeout=5,
        )

    logging.info("Disconnect all clients")
    await proxy.disconnect_all()

    assert len(shared) == 5, shared

    for x in range(5):
        await direct_channel.default_exchange.publish(
            Message(b"1234567890", message_id=f"1-{x}"), queue_name,
        )

    await asyncio.wait_for(reconnect_event.wait(), timeout=5)

    logging.info("Waiting connections")
    await channel.connection.ready()

    async with shared_condition:
        await asyncio.wait_for(
            shared_condition.wait_for(lambda: len(shared) == 10),
            timeout=5,
        )

    assert len(shared) == 10

@aiomisc.timeout(10)
async def test_channel_reconnect(
    connection_fabric, loop, amqp_url, proxy: TCPProxy, add_cleanup: Callable,
):
    heartbeat = 2
    amqp_url = amqp_url.update_query(heartbeat=heartbeat)

    on_reconnect = asyncio.Event()

    conn = await connection_fabric(amqp_url, loop=loop)
    assert isinstance(conn, aio_pika.RobustConnection)

    conn.reconnect_callbacks.add(lambda *_: on_reconnect.set(), weak=False)

    async with conn:
        channel = await conn.channel()
        assert isinstance(channel, aio_pika.RobustChannel)

        async with channel:
            await channel.set_qos(0)
            await channel.set_qos(1)

            with pytest.raises(asyncio.TimeoutError):
                with proxy.slowdown(1, 1):
                    await channel.set_qos(0, timeout=0.5)

            await on_reconnect.wait()
            await channel.set_qos(0)
            await channel.set_qos(1)

@aiomisc.timeout(10)
async def test_channel_reconnect_after_5kb(
    connection: aio_pika.RobustConnection,
    direct_connection: aio_pika.Connection,
    loop: asyncio.AbstractEventLoop,
    proxy: TCPProxy,
    add_cleanup: Callable,
):
    on_reconnect = asyncio.Event()

    connection.reconnect_callbacks.add(
        lambda *_: on_reconnect.set(), weak=False,
    )

    num_bytes = 0

    def write_processor(chunk: bytes) -> bytes:
        nonlocal num_bytes

        if num_bytes >= 0:
            num_bytes += len(chunk)

            if num_bytes < 5000:
                return chunk

        num_bytes = -1
        loop.call_soon(proxy.disconnect_all)
        return chunk

    proxy.write_processor = write_processor

    channel = await connection.channel()
    queue = await channel.declare_queue(auto_delete=False)

    async with direct_connection.channel() as channel:
        for _ in range(10):
            await channel.default_exchange.publish(
                aio_pika.Message(body=b"Hello world"),
                routing_key=queue.name,
            )

    messages = []
    async for message in queue.iterator():
        if len(messages) == 10:
            break
        messages.append(message)

    assert messages
<|endoftext|>


<commit_before>import math
import unittest

try:
    from nose_parameterized import parameterized
except ImportError:
    print("*** Please install 'nose_parameterized' to run these tests ***")
    exit(0)

import oandapyV20.contrib.factories as req

class TestContribFactories(unittest.TestCase):
    """Tests regarding contrib factories.

    The reference is created using the second dict parameter. The
    first dict parameter is merge with this, but only for the keys
    that do NOT exist. That allows us to override parameters.
    The result should reflect the data constructed by the class

    """

    @parameterized.expand([
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {},
           {"len": 1},
        ),
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {"from": "2017-01-01T00:00:00Z",
            "to": "2017-01-02T00:00:00Z",
            "granularity": "M1"},
           {"len": int(math.ceil(24*60.0 / 500))},
        ),
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {"from": "2017-01-01T00:00:00Z",
            "granularity": "M1"},
           {},
        ),
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {"from": "2017-01-01T00:00:00Z",
            "to": "2022-06-30T00:00:00Z",
            "granularity": "M1"},
           {},
        ),
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {"to": "2017-06-30T00:00:00Z",
            "granularity": "M1"},
           {},
           ValueError,
        ),
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {"from": "2017-01-01T00:00:00Z",
            "to": "2017-06-30T00:00:00Z",
            "granularity": "H4"},
           {"len": 3},
        ),
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {"from": "2017-01-01T00:00:00Z",
            "to": "2017-06-30T00:00:00Z",
            "count": 5000,         # same as previous, but increase batchsize
            "granularity": "H4"},
           {"len": 1},
        ),
    ])
    def test__candlehistory(self, factory, instrument, inpar, refpar,
                            exc=None):
        """candlehistoryfactory."""
        if not exc:
            # run the factory
            i = 0
            for r in factory(instrument, params=inpar):
                if i == 0 and inpar:
                    self.assertTrue(r.params['from'] == inpar['from'])
                    # the calculated 'to' should be there
                    self.assertTrue('to' in r.params)
                if 'len' in refpar and i == refpar['len']:
                    self.assertTrue('to' not in r.params)
                i += 1

        # else:
        #     with self.assertRaises(exc) as err:
        #         r = factory(instrument, params=inpar)

if __name__ == "__main__":

    unittest.main()
<commit_msg>extra test<commit_after>import math
import unittest

try:
    from nose_parameterized import parameterized
except ImportError:
    print("*** Please install 'nose_parameterized' to run these tests ***")
    exit(0)

import oandapyV20.contrib.factories as req

class TestContribFactories(unittest.TestCase):
    """Tests regarding contrib factories.

    The reference is created using the second dict parameter. The
    first dict parameter is merge with this, but only for the keys
    that do NOT exist. That allows us to override parameters.
    The result should reflect the data constructed by the class

    """

    @parameterized.expand([
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {},
           {"len": 1},
        ),
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {"from": "2017-01-01T00:00:00Z",
            "to": "2017-01-02T00:00:00Z",
            "granularity": "M1"},
           {"len": int(math.ceil(24*60.0 / 500))},
        ),
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {"from": "2017-01-01T00:00:00Z",
            "granularity": "M1"},
           {},
        ),
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {"from": "2017-01-01T00:00:00Z",
            "to": "2022-06-30T00:00:00Z",
            "granularity": "M1"},
           {},
        ),
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {"to": "2017-06-30T00:00:00Z",
            "granularity": "M1"},
           {},
           ValueError,
        ),
       (req.InstrumentsCandlesFactory,
           "EUR_GBP",
           {"from": "2017-04-01T00:00:00Z",
            "to": "2017-09-23T00:00:00Z",
            "granularity": "M30"},
           {"len": 17},
        ),
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {"from": "2017-01-01T00:00:00Z",
            "to": "2017-06-30T00:00:00Z",
            "granularity": "H4"},
           {"len": 3},
        ),
       (req.InstrumentsCandlesFactory,
           "DE30_EUR",
           {"from": "2017-01-01T00:00:00Z",
            "to": "2017-06-30T00:00:00Z",
            "count": 5000,         # same as previous, but increase batchsize
            "granularity": "H4"},
           {"len": 1},
        ),
    ])
    def test__candlehistory(self, factory, instrument, inpar, refpar,
                            exc=None):
        """candlehistoryfactory."""
        if not exc:
            # run the factory
            i = 0
            for r in factory(instrument, params=inpar):
                if i == 0 and inpar:
                    self.assertTrue(r.params['from'] == inpar['from'])
                    # the calculated 'to' should be there
                    self.assertTrue('to' in r.params)
                if 'len' in refpar and i == refpar['len']:
                    self.assertTrue('to' not in r.params)
                i += 1

        # else:
        #     with self.assertRaises(exc) as err:
        #         r = factory(instrument, params=inpar)

if __name__ == "__main__":

    unittest.main()
<|endoftext|>


<commit_before>
import hashlib
import os.path
import pytest
import json

from ecoli_serotyping.src.ecvalidatingfiles import *

SCRIPT_DIRECTORY = os.path.dirname(os.path.abspath(__file__)) + "/"
REL_DIR = SCRIPT_DIRECTORY + "../Data/Testing_Data/Reference_Genomes"

expectedList1 = sorted([
    REL_DIR + '/NC_011749.1.fasta'
    , REL_DIR + '/NC_011739.1.fasta'
    , REL_DIR + '/NC_011751.1.fasta'
    , REL_DIR + '/AAJT02.1.fsa_nt'
    , REL_DIR + '/NC_002695.1.fasta'
    , REL_DIR + '/NC_002128.1.fasta'
    , REL_DIR + '/NC_011750.1.fasta'
    , REL_DIR + '/NC_002127.1.fasta'])

expectedList2 = sorted([
    SCRIPT_DIRECTORY + '../Data/Testing_Data/Invalid Files/Invalid-File-1.txt',
    SCRIPT_DIRECTORY + '../Data/Testing_Data/Invalid Files/Rabbit-in-Hat.jpg'
])

REL_DIR2 = SCRIPT_DIRECTORY + '../temp/xml/'

expectedList3 = sorted([os.path.abspath(REL_DIR2 + 'AAJT02.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_002127.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_002128.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_002695.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_011739.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_011749.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_011750.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_011751.1.xml')])

def test_getFilesList():

    os.mkdir(SCRIPT_DIRECTORY +'../Data/Testing_Data/Empty_Folder/')

    assert getFilesList(REL_DIR) == expectedList1

    assert getFilesList(REL_DIR + '/NC_011750.1.fasta') == [os.path.abspath(REL_DIR + '/NC_011750.1.fasta')]

    expectedList2.extend(expectedList1)

    assert getFilesList(SCRIPT_DIRECTORY +'../Data/Testing_Data/') == sorted(expectedList2)

    assert getFilesList(SCRIPT_DIRECTORY +'../Data/Testing_Data/Empty_Folder/') == []

def test_getGenomeName():

    assert getGenomeName('AAJT0203.1|gi|190904618|gb|AAJT02000002.1|emb|AAJT0201|dbj|AAJT0202|lcl|AAJT0204|ref|AA_JT0205| Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AAJT0204'

    assert getGenomeName('AAJT0203.1|gi|190904618|gb|AAJT02000002.1|emb|AAJT0201|dbj|AAJT0202|ref|AA_JT0205| Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AAJT0203.1'

    assert getGenomeName('gi|190904618|gb|AAJT02000002.1|emb|AAJT0201|dbj|AAJT0202|ref|AA_JT0205| Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AAJT0200'

    assert getGenomeName('gi|190904618|emb|AAJT0201|dbj|AAJT0202|ref|AA_JT0205 Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AAJT0201'

    assert getGenomeName('gi|190904618|dbj|AAJT0202|ref|AA_JT0205| Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AAJT0202'

    assert getGenomeName('gi|190904618|ref|AA_JT0205| Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AA_JT0205'

    assert getGenomeName('gi|190904618| Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02' ) == '19090461'

    assert getGenomeName('Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AAJT02'

def test_checkFiles():

    assert checkFiles(getFilesList(REL_DIR)) == expectedList1

    assert checkFiles(getFilesList(SCRIPT_DIRECTORY + '../Data/Testing_Data/Invalid Files/')) == 'Error'

    assert checkFiles(getFilesList(SCRIPT_DIRECTORY + '../Data/Testing_Data')) == expectedList1

def test_initializeDB():

    assert subprocess.call(["/usr/bin/makeblastdb", "-in", SCRIPT_DIRECTORY + "../Data/EcOH.fasta ", "-dbtype", "nucl", "-title", "ECTyperDB"]) == 1

    assert subprocess.call(["/usr/bin/makeblastdb", "-in", SCRIPT_DIRECTORY + "../Data/EcOH.fasta ", "-dbtype", "-title", "ECTyperDB", "-out", REL_DIR + "ECTyperDB"]) == 1

    assert subprocess.call(["/usr/bin/makeblastdb", "-in", SCRIPT_DIRECTORY , "-dbtype", "nucl", "-title", "ECTyperDB", "-out", REL_DIR + "ECTyperDB"]) == 1

    ### Tests that won't work outside workspace (to be fixed) ###
    # os.chdir('/home/calarose/Documents/')
    # assert initializeDB() == 0
    #
    # os.chdir('/home/calarose/ectyper/Data/')
    # assert initializeDB() == 0
    #
    # os.chdir('/home/calarose/ectyper/')
    # assert initializeDB() == 0

def test_runBlastQuery():
   #subprocess.call(["mkdir", SCRIPT_DIRECTORY+ "../temp/xml/"])

   os.mkdir(SCRIPT_DIRECTORY+ "../temp/")
   initializeDB()
   assert runBlastQuery(expectedList1, 'ECTyperDB') == expectedList3

def test_parseResults():

    with open(SCRIPT_DIRECTORY + '../Data/Test_dictionaries/ecvalidatingfiles_dict.json') as f:
        expectedDict = json.load(f)

    testDict = parseResults([REL_DIR2 + 'NC_011749.1.xml',REL_DIR2 + 'AAJT02.1.xml'])

    for test_genome, test_alignment in testDict.iteritems():
        if test_genome in expectedDict:
            if isinstance(test_alignment, dict):
                for test_title, test_info in test_alignment.iteritems():
                    test_length = test_info.keys()[0]
                    test_hsp = test_info[test_length]
                    hash_hsp = hashlib.md5()
                    hash_hsp.update(str(test_hsp))
                    assert expectedDict[test_genome][test_title] == str(hash_hsp.hexdigest())
                    # if expectedDict[test_genome][test_title] == str(hash_hsp.hexdigest()):
                    #     pytest.fail("The dictionaries aren't the same. Test failed. \nCAUSE \nAlignment title: " + str(test_title)  + '\nExpected: ' + str(hash_hsp.hexdigest()) +  '\nFound: ' +
                    #                 str(expectedDict[test_genome][test_title]) )
        else:
            pytest.fail("The dictionaries aren't the same. Test failed. \nCAUSE \n" + str(test_genome) + " isn't a key")

<commit_msg>Finalized tests and creation of folders if needed in tests.<commit_after>
import hashlib
import os.path
import pytest
import json

from ecoli_serotyping.src.ecvalidatingfiles import *

SCRIPT_DIRECTORY = os.path.dirname(os.path.abspath(__file__)) + "/"
REL_DIR = SCRIPT_DIRECTORY + "../Data/Testing_Data/Reference_Genomes"

expectedList1 = sorted([
    REL_DIR + '/NC_011749.1.fasta'
    , REL_DIR + '/NC_011739.1.fasta'
    , REL_DIR + '/NC_011751.1.fasta'
    , REL_DIR + '/AAJT02.1.fsa_nt'
    , REL_DIR + '/NC_002695.1.fasta'
    , REL_DIR + '/NC_002128.1.fasta'
    , REL_DIR + '/NC_011750.1.fasta'
    , REL_DIR + '/NC_002127.1.fasta'])

expectedList2 = sorted([
    SCRIPT_DIRECTORY + '../Data/Testing_Data/Invalid Files/Invalid-File-1.txt',
    SCRIPT_DIRECTORY + '../Data/Testing_Data/Invalid Files/Rabbit-in-Hat.jpg'
])

REL_DIR2 = SCRIPT_DIRECTORY + '../temp/xml/'

expectedList3 = sorted([os.path.abspath(REL_DIR2 + 'AAJT02.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_002127.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_002128.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_002695.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_011739.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_011749.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_011750.1.xml'),
                        os.path.abspath(REL_DIR2 + 'NC_011751.1.xml')])

def test_getFilesList():

    if not os.path.isdir(SCRIPT_DIRECTORY +'../Data/Testing_Data/Empty_Folder/'):
        os.mkdir(SCRIPT_DIRECTORY +'../Data/Testing_Data/Empty_Folder/')

    assert getFilesList(REL_DIR) == expectedList1

    assert getFilesList(REL_DIR + '/NC_011750.1.fasta') == [os.path.abspath(REL_DIR + '/NC_011750.1.fasta')]

    expectedList2.extend(expectedList1)

    assert getFilesList(SCRIPT_DIRECTORY +'../Data/Testing_Data/') == sorted(expectedList2)

    assert getFilesList(SCRIPT_DIRECTORY +'../Data/Testing_Data/Empty_Folder/') == []

def test_getGenomeName():

    assert getGenomeName('AAJT0203.1|gi|190904618|gb|AAJT02000002.1|emb|AAJT0201|dbj|AAJT0202|lcl|AAJT0204|ref|AA_JT0205| Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AAJT0204'

    assert getGenomeName('AAJT0203.1|gi|190904618|gb|AAJT02000002.1|emb|AAJT0201|dbj|AAJT0202|ref|AA_JT0205| Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AAJT0203.1'

    assert getGenomeName('gi|190904618|gb|AAJT02000002.1|emb|AAJT0201|dbj|AAJT0202|ref|AA_JT0205| Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AAJT0200'

    assert getGenomeName('gi|190904618|emb|AAJT0201|dbj|AAJT0202|ref|AA_JT0205 Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AAJT0201'

    assert getGenomeName('gi|190904618|dbj|AAJT0202|ref|AA_JT0205| Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AAJT0202'

    assert getGenomeName('gi|190904618|ref|AA_JT0205| Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AA_JT0205'

    assert getGenomeName('gi|190904618| Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02' ) == '19090461'

    assert getGenomeName('Escherichia coli B7A gcontig_1112495745514, whole genome shotgun sequence', 'AAJT02') == 'AAJT02'

def test_checkFiles():

    assert checkFiles(getFilesList(REL_DIR)) == expectedList1

    assert checkFiles(getFilesList(SCRIPT_DIRECTORY + '../Data/Testing_Data/Invalid Files/')) == 'Error'

    assert checkFiles(getFilesList(SCRIPT_DIRECTORY + '../Data/Testing_Data')) == expectedList1

def test_initializeDB():

    assert subprocess.call(["/usr/bin/makeblastdb", "-in", SCRIPT_DIRECTORY + "../Data/EcOH.fasta ", "-dbtype", "nucl", "-title", "ECTyperDB"]) == 1

    assert subprocess.call(["/usr/bin/makeblastdb", "-in", SCRIPT_DIRECTORY + "../Data/EcOH.fasta ", "-dbtype", "-title", "ECTyperDB", "-out", REL_DIR + "ECTyperDB"]) == 1

    assert subprocess.call(["/usr/bin/makeblastdb", "-in", SCRIPT_DIRECTORY , "-dbtype", "nucl", "-title", "ECTyperDB", "-out", REL_DIR + "ECTyperDB"]) == 1

    ### Tests that won't work outside workspace (to be fixed) ###
    # os.chdir('/home/calarose/Documents/')
    # assert initializeDB() == 0
    #
    # os.chdir('/home/calarose/ectyper/Data/')
    # assert initializeDB() == 0
    #
    # os.chdir('/home/calarose/ectyper/')
    # assert initializeDB() == 0

def test_runBlastQuery():
   #subprocess.call(["mkdir", SCRIPT_DIRECTORY+ "../temp/xml/"])
   if not os.path.isdir(SCRIPT_DIRECTORY+ "../temp/"):
        os.mkdir(SCRIPT_DIRECTORY+ "../temp/")
   initializeDB()
   assert runBlastQuery(expectedList1, 'ECTyperDB') == expectedList3

def test_parseResults():

    with open(SCRIPT_DIRECTORY + '../Data/Test_dictionaries/ecvalidatingfiles_dict.json') as f:
        expectedDict = json.load(f)

    testDict = parseResults([REL_DIR2 + 'NC_011749.1.xml',REL_DIR2 + 'AAJT02.1.xml'])

    for test_genome, test_alignment in testDict.iteritems():
        if test_genome in expectedDict:
            if isinstance(test_alignment, dict):
                for test_title, test_info in test_alignment.iteritems():
                    test_length = test_info.keys()[0]
                    test_hsp = test_info[test_length]
                    hash_hsp = hashlib.md5()
                    hash_hsp.update(str(test_hsp))
                    assert expectedDict[test_genome][test_title] == str(hash_hsp.hexdigest())
                    # if expectedDict[test_genome][test_title] == str(hash_hsp.hexdigest()):
                    #     pytest.fail("The dictionaries aren't the same. Test failed. \nCAUSE \nAlignment title: " + str(test_title)  + '\nExpected: ' + str(hash_hsp.hexdigest()) +  '\nFound: ' +
                    #                 str(expectedDict[test_genome][test_title]) )
        else:
            pytest.fail("The dictionaries aren't the same. Test failed. \nCAUSE \n" + str(test_genome) + " isn't a key")

<|endoftext|>


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

#   Copyright (c) 2010-2016, MIT Probabilistic Computing Project
#
#   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.

import collections
import itertools
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
import pytest
import random
from stochastic import stochastic
import struct
from scipy import stats
import tempfile
import time

import bayeslite
from bayeslite import read_pandas
from bayeslite.backends.loom_backend import LoomBackend

def axis_aligned_gaussians(means, size, seed=None):
    if seed is not None:
        np.random.seed(seed)
    return [np.random.multivariate_normal(mean, [[0,1], [1,0]], size=size)
            for mean
            in means]

def mix(gaussians, p, seed=None):
    assert len(gaussians) == len(p)
    assert reduce(lambda sum, n: sum + n, p) == 1

    if seed is not None:
        np.random.seed(seed)
        choices = np.random.choice([0,1], size=len(gaussians[0]), p=p)
    return [gaussians[choices[i]][i] for i in range(len(choices))]

def temp_file_path(suffix):
    """Returns a random file path within '/tmp'."""
    _, temp_filename = tempfile.mkstemp(suffix=suffix)
    os.remove(temp_filename)
    return temp_filename

def register_loom(bdb):
    loom_store_path = temp_file_path('.bdb')
    loom_backend = LoomBackend(loom_store_path=loom_store_path)
    bayeslite.bayesdb_register_backend(bdb, loom_backend)

def distance(a, b):
    '''
    Computes the euclidian distance between points `a` and `b`.
    '''
    # coerce to `np.ndarray`s if necessary
    a = a if isinstance(a, np.ndarray) else np.array(a)
    b = b if isinstance(a, np.ndarray) else np.array(b)
    return abs(np.linalg.norm(a-b))

def prepare_bdb(bdb, samples, table, seed):
    quoted_table = bayeslite.bql_quote_name(table)
    dataframe = pd.DataFrame(data=samples)
    read_pandas.bayesdb_read_pandas_df(bdb, 'data', dataframe, create=True)

    bdb.execute('''
        CREATE POPULATION FOR %s WITH SCHEMA (
            GUESS STATTYPES OF (*)
        )
    ''' % (quoted_table,))
    bdb.execute('CREATE GENERATOR FOR %s USING loom;' % (quoted_table,))
    bdb.execute('INITIALIZE 4 MODELS FOR %s;' % (quoted_table,))
    bdb.execute('ANALYZE %s FOR 100 ITERATIONS;' % (quoted_table,))

def insert_row(bdb, table, x, y):
    quoted_table = bayeslite.bql_quote_name(table)
    query = 'INSERT INTO {table} ("0", "1") VALUES (?, ?)'.format(table=quoted_table)
    bdb.sql_execute(query, bindings=(x, y))
    return bdb.sql_execute('SELECT last_insert_rowid()').fetchone()[0]

def simulate_from_rowid(bdb, table, column, rowid, limit=1000):
    quoted_table = bayeslite.bql_quote_name(table)
    quoted_column = bayeslite.bql_quote_name(str(column))
    query = 'SIMULATE {column} FROM {table} GIVEN rowid=? LIMIT ?'.format(
        column=quoted_column,
        table=quoted_table
    )
    cursor = bdb.execute(query, (rowid, limit))
    return [float(x[0]) for x in cursor]

@stochastic(max_runs=1, min_passes=1)
def test_mix_ratio(seed):
    means = ((0,20), (20,0))
    sample_size = 100
    mix_ratio = [0.7, 0.3]
    table = 'data'

    sample_gaussians = axis_aligned_gaussians(means, sample_size, seed=seed)
    samples = mix(sample_gaussians, mix_ratio, seed=seed)

    packed_seed = struct.pack('<QQQQ', seed, seed, seed, seed)
    with bayeslite.bayesdb_open(seed=packed_seed) as bdb:
        register_loom(bdb)
        prepare_bdb(bdb, samples, table, seed)

        cursor = bdb.execute('SIMULATE "0", "1" FROM data LIMIT ?;', (sample_size,))
        simulated_samples = [sample for sample in cursor]

    counts = collections.Counter(
        (0 if distance((x,y), means[0]) < distance((x,y), means[1]) else 1
         for x, y
         in simulated_samples)
    )
    simulated_mix_ratio = [counts[key] / float(len(simulated_samples)) for key in counts]

    for i in range(len(means)):
        difference = abs(mix_ratio[i] - simulated_mix_ratio[i])
        assert difference < 0.1

@stochastic(max_runs=1, min_passes=1)
def test_simulate_y_from_partially_populated_row(seed):
    means = ((0,20), (20,0))
    sample_size = 50
    mix_ratio = [0.7, 0.3]
    table = 'data'

    sample_gaussians = axis_aligned_gaussians(means, sample_size, seed=seed)
    samples = mix(sample_gaussians, mix_ratio, seed=seed)

    packed_seed = struct.pack('<QQQQ', seed, seed, seed, seed)
    with bayeslite.bayesdb_open(seed=packed_seed) as bdb:
        register_loom(bdb)
        prepare_bdb(bdb, samples, table, seed)

        rowid = insert_row(bdb, table, means[0][0], None)
        simulated_samples = simulate_from_rowid(bdb, table, 1, rowid, limit=sample_size)

    y_samples = [y for x, y in sample_gaussians[0]]
    statistic, p_value = stats.ks_2samp(y_samples, simulated_samples)
    assert(statistic < 0.1 or p_value > 0.01)

def test_simulate_conflict():
    with bayeslite.bayesdb_open() as bdb:
        bdb.sql_execute('''
            CREATE TABLE data (
                '0' numeric PRIMARY KEY,
                '1' numeric
            );
        ''')
        insert_row(bdb, 'data', 1, 1)
        bdb.execute('''
            CREATE POPULATION FOR data WITH SCHEMA (
                "0" NUMERICAL;
                "1" NUMERICAL;
            );
        ''')
        bdb.execute('CREATE GENERATOR FOR data USING cgpm;')
        bdb.execute('INITIALIZE 1 MODELS FOR data;')

        rowid = insert_row(bdb, 'data', 0, None)
        with pytest.raises(Exception):
            bdb.execute('''
                SIMULATE "0" FROM data
                    GIVEN rowid=?, "0"= 0, 1"=0
                    LIMIT 1;
            ''', (rowid,))
<commit_msg>Fix use of stochastic decorator<commit_after># -*- coding: utf-8 -*-

#   Copyright (c) 2010-2016, MIT Probabilistic Computing Project
#
#   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.

import collections
import itertools
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
import pytest
import random
from stochastic import stochastic
import struct
from scipy import stats
import tempfile
import time

import bayeslite
from bayeslite import read_pandas
from bayeslite.backends.loom_backend import LoomBackend

def set_numpy_seed(seed):
    """
    The `stochastic` decorator provides 64 byte strings as its seeds. numpy
    uses 32 bit unsigned integers as its seeds Takes a 64 byte seed, extracts
    the first 32 bits (4 bytes) from it and uses those as the numpy random
    seed.
    """
    np_seed = struct.unpack("!I", seed[:4])[0] # convert to 32 bit uint
    np.random.seed(np_seed)

def axis_aligned_gaussians(means, size, seed=None):
    if seed is not None:
        set_numpy_seed(seed)
    return [np.random.multivariate_normal(mean, [[0,1], [1,0]], size=size)
            for mean
            in means]

def mix(gaussians, p, seed=None):
    assert len(gaussians) == len(p)
    assert reduce(lambda sum, n: sum + n, p) == 1

    if seed is not None:
        set_numpy_seed(seed)
        choices = np.random.choice([0,1], size=len(gaussians[0]), p=p)
    return [gaussians[choices[i]][i] for i in range(len(choices))]

def temp_file_path(suffix):
    """Returns a random file path within '/tmp'."""
    _, temp_filename = tempfile.mkstemp(suffix=suffix)
    os.remove(temp_filename)
    return temp_filename

def register_loom(bdb):
    loom_store_path = temp_file_path('.bdb')
    loom_backend = LoomBackend(loom_store_path=loom_store_path)
    bayeslite.bayesdb_register_backend(bdb, loom_backend)

def distance(a, b):
    '''
    Computes the euclidian distance between points `a` and `b`.
    '''
    # coerce to `np.ndarray`s if necessary
    a = a if isinstance(a, np.ndarray) else np.array(a)
    b = b if isinstance(a, np.ndarray) else np.array(b)
    return abs(np.linalg.norm(a-b))

def prepare_bdb(bdb, samples, table, seed):
    quoted_table = bayeslite.bql_quote_name(table)
    dataframe = pd.DataFrame(data=samples)
    read_pandas.bayesdb_read_pandas_df(bdb, 'data', dataframe, create=True)

    bdb.execute('''
        CREATE POPULATION FOR %s WITH SCHEMA (
            GUESS STATTYPES OF (*)
        )
    ''' % (quoted_table,))
    bdb.execute('CREATE GENERATOR FOR %s USING loom;' % (quoted_table,))
    bdb.execute('INITIALIZE 4 MODELS FOR %s;' % (quoted_table,))
    bdb.execute('ANALYZE %s FOR 100 ITERATIONS;' % (quoted_table,))

def insert_row(bdb, table, x, y):
    quoted_table = bayeslite.bql_quote_name(table)
    query = 'INSERT INTO {table} ("0", "1") VALUES (?, ?)'.format(table=quoted_table)
    bdb.sql_execute(query, bindings=(x, y))
    return bdb.sql_execute('SELECT last_insert_rowid()').fetchone()[0]

def simulate_from_rowid(bdb, table, column, rowid, limit=1000):
    quoted_table = bayeslite.bql_quote_name(table)
    quoted_column = bayeslite.bql_quote_name(str(column))
    query = 'SIMULATE {column} FROM {table} GIVEN rowid=? LIMIT ?'.format(
        column=quoted_column,
        table=quoted_table
    )
    cursor = bdb.execute(query, (rowid, limit))
    return [float(x[0]) for x in cursor]

@stochastic(max_runs=1, min_passes=1)
def test_mix_ratio(seed):
    means = ((0,20), (20,0))
    sample_size = 100
    mix_ratio = [0.7, 0.3]
    table = 'data'

    sample_gaussians = axis_aligned_gaussians(means, sample_size, seed=seed)
    samples = mix(sample_gaussians, mix_ratio, seed=seed)

    with bayeslite.bayesdb_open(seed=seed) as bdb:
        register_loom(bdb)
        prepare_bdb(bdb, samples, table, seed)

        cursor = bdb.execute('SIMULATE "0", "1" FROM data LIMIT ?;', (sample_size,))
        simulated_samples = [sample for sample in cursor]

    counts = collections.Counter(
        (0 if distance((x,y), means[0]) < distance((x,y), means[1]) else 1
         for x, y
         in simulated_samples)
    )
    simulated_mix_ratio = [counts[key] / float(len(simulated_samples)) for key in counts]

    for i in range(len(means)):
        difference = abs(mix_ratio[i] - simulated_mix_ratio[i])
        assert difference < 0.1

@stochastic(max_runs=1, min_passes=1)
def test_simulate_y_from_partially_populated_row(seed):
    means = ((0,20), (20,0))
    sample_size = 50
    mix_ratio = [0.7, 0.3]
    table = 'data'

    sample_gaussians = axis_aligned_gaussians(means, sample_size, seed=seed)
    samples = mix(sample_gaussians, mix_ratio, seed=seed)

    with bayeslite.bayesdb_open(seed=seed) as bdb:
        register_loom(bdb)
        prepare_bdb(bdb, samples, table, seed)

        rowid = insert_row(bdb, table, means[0][0], None)
        simulated_samples = simulate_from_rowid(bdb, table, 1, rowid, limit=sample_size)

    y_samples = [y for x, y in sample_gaussians[0]]
    statistic, p_value = stats.ks_2samp(y_samples, simulated_samples)
    assert(statistic < 0.1 or p_value > 0.01)

def test_simulate_conflict():
    with bayeslite.bayesdb_open() as bdb:
        bdb.sql_execute('''
            CREATE TABLE data (
                '0' numeric PRIMARY KEY,
                '1' numeric
            );
        ''')
        insert_row(bdb, 'data', 1, 1)
        bdb.execute('''
            CREATE POPULATION FOR data WITH SCHEMA (
                "0" NUMERICAL;
                "1" NUMERICAL;
            );
        ''')
        bdb.execute('CREATE GENERATOR FOR data USING cgpm;')
        bdb.execute('INITIALIZE 1 MODELS FOR data;')

        rowid = insert_row(bdb, 'data', 0, None)
        with pytest.raises(Exception):
            bdb.execute('''
                SIMULATE "0" FROM data
                    GIVEN rowid=?, "0"= 0, 1"=0
                    LIMIT 1;
            ''', (rowid,))
<|endoftext|>


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

#   Copyright (c) 2010-2016, MIT Probabilistic Computing Project
#
#   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.

import collections
import os
import struct
import tempfile

import numpy as np
import pandas as pd
import pytest

from scipy import stats

import bayeslite

from bayeslite import read_pandas
from bayeslite.backends.loom_backend import LoomBackend

from stochastic import stochastic

def axis_aligned_gaussians(means, size, rng):
    return [rng.multivariate_normal(mean, [[0,1], [1,0]], size=size)
            for mean in means]

def mix(gaussians, p, rng):
    assert len(gaussians) == len(p)
    assert reduce(lambda sum, n: sum + n, p) == 1

    choices = rng.choice([0,1], size=len(gaussians[0]), p=p)
    return [gaussians[choices[i]][i] for i in range(len(choices))]

def temp_file_path(suffix):
    """Returns a random file path within '/tmp'."""
    _, temp_filename = tempfile.mkstemp(suffix=suffix)
    os.remove(temp_filename)
    return temp_filename

def register_loom(bdb):
    loom_store_path = temp_file_path('.bdb')
    loom_backend = LoomBackend(loom_store_path=loom_store_path)
    bayeslite.bayesdb_register_backend(bdb, loom_backend)

def distance(a, b):
    """Computes the Euclidean distance between points `a` and `b`."""
    a = a if isinstance(a, np.ndarray) else np.array(a)
    b = b if isinstance(a, np.ndarray) else np.array(b)
    return abs(np.linalg.norm(a-b))

def prepare_bdb(bdb, samples, table):
    quoted_table = bayeslite.bql_quote_name(table)
    dataframe = pd.DataFrame(data=samples)
    read_pandas.bayesdb_read_pandas_df(bdb, 'data', dataframe, create=True)

    bdb.execute('''
        CREATE POPULATION FOR %s WITH SCHEMA (
            GUESS STATTYPES OF (*)
        )
    ''' % (quoted_table,))
    bdb.execute('CREATE GENERATOR FOR %s USING loom;' % (quoted_table,))
    bdb.execute('INITIALIZE 4 MODELS FOR %s;' % (quoted_table,))
    bdb.execute('ANALYZE %s FOR 100 ITERATIONS;' % (quoted_table,))

def insert_row(bdb, table, x, y):
    quoted_table = bayeslite.bql_quote_name(table)
    query = 'INSERT INTO {table} ("0", "1") VALUES (?, ?)'.format(table=quoted_table)
    bdb.sql_execute(query, bindings=(x, y))
    return bdb.sql_execute('SELECT last_insert_rowid()').fetchone()[0]

def simulate_from_rowid(bdb, table, column, rowid, limit=1000):
    quoted_table = bayeslite.bql_quote_name(table)
    quoted_column = bayeslite.bql_quote_name(str(column))
    query = 'SIMULATE {column} FROM {table} GIVEN rowid=? LIMIT ?'.format(
        column=quoted_column,
        table=quoted_table
    )
    cursor = bdb.execute(query, (rowid, limit))
    return [float(x[0]) for x in cursor]

@stochastic(max_runs=1, min_passes=1)
def test_mix_ratio(seed):
    means = ((0,20), (20,0))
    sample_size = 100
    mix_ratio = [0.7, 0.3]
    table = 'data'

    with bayeslite.bayesdb_open(seed=seed) as bdb:
        sample_gaussians = axis_aligned_gaussians(means, sample_size, bdb._np_prng)
        samples = mix(sample_gaussians, mix_ratio, bdb._np_prng)
        register_loom(bdb)
        prepare_bdb(bdb, samples, table)

        cursor = bdb.execute('SIMULATE "0", "1" FROM data LIMIT ?;', (sample_size,))
        simulated_samples = [sample for sample in cursor]

    counts = collections.Counter(
        (0 if distance((x,y), means[0]) < distance((x,y), means[1]) else 1
         for x, y
         in simulated_samples)
    )
    simulated_mix_ratio = [counts[key] / float(len(simulated_samples)) for key in counts]

    for i in range(len(means)):
        difference = abs(mix_ratio[i] - simulated_mix_ratio[i])
        assert difference < 0.1

@stochastic(max_runs=1, min_passes=1)
def test_simulate_y_from_partially_populated_row(seed):
    means = ((0,20), (20,0))
    sample_size = 50
    mix_ratio = [0.7, 0.3]
    table = 'data'

    with bayeslite.bayesdb_open(seed=seed) as bdb:
        sample_gaussians = axis_aligned_gaussians(means, sample_size, bdb._np_prng)
        samples = mix(sample_gaussians, mix_ratio, bdb._np_prng)
        register_loom(bdb)
        prepare_bdb(bdb, samples, table)

        rowid = insert_row(bdb, table, means[0][0], None)
        simulated_samples = simulate_from_rowid(bdb, table, 1, rowid, limit=sample_size)

    y_samples = [y for x, y in sample_gaussians[0]]
    statistic, p_value = stats.ks_2samp(y_samples, simulated_samples)
    assert(statistic < 0.1 or p_value > 0.01)

def test_simulate_conflict():
    """Cannot override existing value in table using GIVEN in SIMULATE."""
    with bayeslite.bayesdb_open() as bdb:
        bdb.sql_execute('''
            CREATE TABLE data (
                '0' numeric PRIMARY KEY,
                '1' numeric
            );
        ''')
        insert_row(bdb, 'data', 1, 1)
        bdb.execute('''
            CREATE POPULATION FOR data WITH SCHEMA (
                "0" NUMERICAL;
                "1" NUMERICAL;
            );
        ''')
        bdb.execute('CREATE GENERATOR FOR data USING cgpm;')
        bdb.execute('INITIALIZE 1 MODELS FOR data;')

        rowid = insert_row(bdb, 'data', 0, None)
        with pytest.raises(Exception):
            bdb.execute('''
                SIMULATE "0" FROM data
                    GIVEN rowid=?, "0"= 0, 1"=0
                    LIMIT 1;
            ''', (rowid,))
<commit_msg>Shortened variable names.<commit_after># -*- coding: utf-8 -*-

#   Copyright (c) 2010-2016, MIT Probabilistic Computing Project
#
#   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.

import collections
import os
import struct
import tempfile

import numpy as np
import pandas as pd
import pytest

from scipy import stats

import bayeslite

from bayeslite import read_pandas
from bayeslite.backends.loom_backend import LoomBackend

from stochastic import stochastic

def axis_aligned_gaussians(means, size, rng):
    return [rng.multivariate_normal(mean, [[0,1], [1,0]], size=size)
            for mean in means]

def mix(gaussians, p, rng):
    assert len(gaussians) == len(p)
    assert reduce(lambda sum, n: sum + n, p) == 1

    choices = rng.choice([0,1], size=len(gaussians[0]), p=p)
    return [gaussians[choices[i]][i] for i in range(len(choices))]

def temp_file_path(suffix):
    """Returns a random file path within '/tmp'."""
    _, temp_filename = tempfile.mkstemp(suffix=suffix)
    os.remove(temp_filename)
    return temp_filename

def register_loom(bdb):
    loom_store_path = temp_file_path('.bdb')
    loom_backend = LoomBackend(loom_store_path=loom_store_path)
    bayeslite.bayesdb_register_backend(bdb, loom_backend)

def distance(a, b):
    """Computes the Euclidean distance between points `a` and `b`."""
    a = a if isinstance(a, np.ndarray) else np.array(a)
    b = b if isinstance(a, np.ndarray) else np.array(b)
    return abs(np.linalg.norm(a-b))

def prepare_bdb(bdb, samples, table):
    qt = bayeslite.bql_quote_name(table)
    dataframe = pd.DataFrame(data=samples)
    read_pandas.bayesdb_read_pandas_df(bdb, 'data', dataframe, create=True)

    bdb.execute('''
        CREATE POPULATION FOR %s WITH SCHEMA (
            GUESS STATTYPES OF (*)
        )
    ''' % (qt,))
    bdb.execute('CREATE GENERATOR FOR %s USING loom;' % (qt,))
    bdb.execute('INITIALIZE 4 MODELS FOR %s;' % (qt,))
    bdb.execute('ANALYZE %s FOR 100 ITERATIONS;' % (qt,))

def insert_row(bdb, table, x, y):
    qt = bayeslite.bql_quote_name(table)
    query = 'INSERT INTO %s ("0", "1") VALUES (?, ?)' % (qt,)
    bdb.sql_execute(query, bindings=(x, y))
    return bdb.sql_execute('SELECT last_insert_rowid()').fetchone()[0]

def simulate_from_rowid(bdb, table, column, rowid, limit=1000):
    qt = bayeslite.bql_quote_name(table)
    qc = bayeslite.bql_quote_name(str(column))
    cursor = bdb.execute('''
        SIMULATE %s FROM %s GIVEN rowid=? LIMIT ?
    ''' % (qc, qt) , bindings=(rowid, limit))
    return [float(x[0]) for x in cursor]

@stochastic(max_runs=1, min_passes=1)
def test_mix_ratio(seed):
    means = ((0,20), (20,0))
    sample_size = 100
    mix_ratio = [0.7, 0.3]
    table = 'data'

    with bayeslite.bayesdb_open(seed=seed) as bdb:
        sample_gaussians = axis_aligned_gaussians(means, sample_size, bdb._np_prng)
        samples = mix(sample_gaussians, mix_ratio, bdb._np_prng)
        register_loom(bdb)
        prepare_bdb(bdb, samples, table)

        cursor = bdb.execute('SIMULATE "0", "1" FROM data LIMIT ?;', (sample_size,))
        simulated_samples = [sample for sample in cursor]

    counts = collections.Counter(
        (0 if distance((x,y), means[0]) < distance((x,y), means[1]) else 1
         for x, y
         in simulated_samples)
    )
    simulated_mix_ratio = [counts[key] / float(len(simulated_samples)) for key in counts]

    for i in range(len(means)):
        difference = abs(mix_ratio[i] - simulated_mix_ratio[i])
        assert difference < 0.1

@stochastic(max_runs=1, min_passes=1)
def test_simulate_y_from_partially_populated_row(seed):
    means = ((0,20), (20,0))
    sample_size = 50
    mix_ratio = [0.7, 0.3]
    table = 'data'

    with bayeslite.bayesdb_open(seed=seed) as bdb:
        sample_gaussians = axis_aligned_gaussians(means, sample_size, bdb._np_prng)
        samples = mix(sample_gaussians, mix_ratio, bdb._np_prng)
        register_loom(bdb)
        prepare_bdb(bdb, samples, table)

        rowid = insert_row(bdb, table, means[0][0], None)
        simulated_samples = simulate_from_rowid(bdb, table, 1, rowid, limit=sample_size)

    y_samples = [y for x, y in sample_gaussians[0]]
    statistic, p_value = stats.ks_2samp(y_samples, simulated_samples)
    assert(statistic < 0.1 or p_value > 0.01)

def test_simulate_conflict():
    """Cannot override existing value in table using GIVEN in SIMULATE."""
    with bayeslite.bayesdb_open() as bdb:
        bdb.sql_execute('''
            CREATE TABLE data (
                '0' numeric PRIMARY KEY,
                '1' numeric
            );
        ''')
        insert_row(bdb, 'data', 1, 1)
        bdb.execute('''
            CREATE POPULATION FOR data WITH SCHEMA (
                "0" NUMERICAL;
                "1" NUMERICAL;
            );
        ''')
        bdb.execute('CREATE GENERATOR FOR data USING cgpm;')
        bdb.execute('INITIALIZE 1 MODELS FOR data;')

        rowid = insert_row(bdb, 'data', 0, None)
        with pytest.raises(Exception):
            bdb.execute('''
                SIMULATE "0" FROM data
                    GIVEN rowid=?, "0"= 0, 1"=0
                    LIMIT 1;
            ''', (rowid,))
<|endoftext|>


<commit_before>import unittest
import lattice_mc
import random
import numpy as np

class IntegrationTestCase( unittest.TestCase ):

    def test_simulation_runs_with_selected_sites( self ):
        a, b, c, = 4, 4, 4
        spacing = 1.0
        n_atoms = 32
        selected_sites = 'L'
        number_of_jumps = 10
        s = lattice_mc.Simulation()
        s.lattice = lattice_mc.init_lattice.cubic_lattice( a, b, c, spacing )
        s.lattice.transmute_sites( 'L', 'X', n_atoms )
        s.set_number_of_atoms( n_atoms, selected_sites=selected_sites )
        [ self.assertEqual( site.label, selected_sites ) for site in s.lattice.occupied_sites() ]
        s.setup_lookup_table()
        s.set_number_of_jumps( number_of_jumps )
        self.assertEqual( s.tracer_correlation, None )
        self.assertEqual( s.collective_correlation, None )
        self.assertEqual( s.tracer_diffusion_coefficient, None )
        self.assertEqual( s.collective_diffusion_coefficient_per_atom, None )
        self.assertEqual( s.average_site_occupations, None )
        s.run()
        self.assertNotEqual( s.tracer_correlation, None )
        self.assertNotEqual( s.collective_correlation, None )
        self.assertNotEqual( s.tracer_diffusion_coefficient, None )
        self.assertNotEqual( s.collective_diffusion_coefficient_per_atom, None )
        self.assertNotEqual( s.average_site_occupations, None )

    def test_simulation_runs_with_variable_coordination_numbers( self ):
        s = lattice_mc.Simulation()
        site_data = [ [ 1, np.array( [ 0.0, 0.0, 0.0 ] ), [ 2 ],    0.0, 'A' ],
                      [ 2, np.array( [ 1.0, 0.0, 0.0 ] ), [ 1, 3 ], 0.0, 'A' ],
                      [ 3, np.array( [ 2.0, 1.0, 1.0 ] ), [ 2 ],    0.0, 'A' ] ]
        for d in site_data:
            print( *d )
        sites = [ lattice_mc.lattice_site.Site( *d ) for d in site_data ]
        s.lattice = lattice_mc.lattice.Lattice( sites, cell_lengths=np.array( [ 10.0, 10.0, 10.0 ] ) )
        s.set_number_of_atoms( 1 )
        s.setup_lookup_table()
        s.set_number_of_jumps( 10 )
        s.run()

if __name__ == '__main__':
    unittest.main()
<commit_msg>Removed left-over print() statement<commit_after>import unittest
import lattice_mc
import random
import numpy as np

class IntegrationTestCase( unittest.TestCase ):

    def test_simulation_runs_with_selected_sites( self ):
        a, b, c, = 4, 4, 4
        spacing = 1.0
        n_atoms = 32
        selected_sites = 'L'
        number_of_jumps = 10
        s = lattice_mc.Simulation()
        s.lattice = lattice_mc.init_lattice.cubic_lattice( a, b, c, spacing )
        s.lattice.transmute_sites( 'L', 'X', n_atoms )
        s.set_number_of_atoms( n_atoms, selected_sites=selected_sites )
        [ self.assertEqual( site.label, selected_sites ) for site in s.lattice.occupied_sites() ]
        s.setup_lookup_table()
        s.set_number_of_jumps( number_of_jumps )
        self.assertEqual( s.tracer_correlation, None )
        self.assertEqual( s.collective_correlation, None )
        self.assertEqual( s.tracer_diffusion_coefficient, None )
        self.assertEqual( s.collective_diffusion_coefficient_per_atom, None )
        self.assertEqual( s.average_site_occupations, None )
        s.run()
        self.assertNotEqual( s.tracer_correlation, None )
        self.assertNotEqual( s.collective_correlation, None )
        self.assertNotEqual( s.tracer_diffusion_coefficient, None )
        self.assertNotEqual( s.collective_diffusion_coefficient_per_atom, None )
        self.assertNotEqual( s.average_site_occupations, None )

    def test_simulation_runs_with_variable_coordination_numbers( self ):
        s = lattice_mc.Simulation()
        site_data = [ [ 1, np.array( [ 0.0, 0.0, 0.0 ] ), [ 2 ],    0.0, 'A' ],
                      [ 2, np.array( [ 1.0, 0.0, 0.0 ] ), [ 1, 3 ], 0.0, 'A' ],
                      [ 3, np.array( [ 2.0, 1.0, 1.0 ] ), [ 2 ],    0.0, 'A' ] ]
        sites = [ lattice_mc.lattice_site.Site( *d ) for d in site_data ]
        s.lattice = lattice_mc.lattice.Lattice( sites, cell_lengths=np.array( [ 10.0, 10.0, 10.0 ] ) )
        s.set_number_of_atoms( 1 )
        s.setup_lookup_table()
        s.set_number_of_jumps( 10 )
        s.run()

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