<commit_before>import logging
import pytest
import sys
import time
import zmq

from test import support

from mdbase.broker import Service, Worker, MajorDomoBroker, W_READY

log = logging.getLogger()
log.addHandler(logging.StreamHandler(sys.stdout))

@pytest.fixture
def broker():
    return MajorDomoBroker()

@pytest.fixture
def broker_verbose():
    return MajorDomoBroker(True)

@pytest.fixture
def address():
    return b"tcp://localhost:6666"

@pytest.fixture
def service():
    return Service(b"S_ECHO")

class TestBrokerService():
    def test_service_instantiate(self):
        """Test instantiating service model"""
        name = "Service1"
        s = Service(name)
        assert s.requests == []
        assert s.waiting == []
        assert s.name == name

class TestBrokerWorker():
    def test_worker_instantiate(self):
        """Test instantiating worker model"""
        identity = hex(1)
        address = "worker1"
        expiry = 3
        w = Worker(identity, address, expiry)
        assert w.service is None
        assert w.identity == identity
        assert w.address == address
        assert w.expiry is not None

class TestBrokerModel():
    def test_broker_instantiate(self):
        """Test instantiating majordomo worker model"""
        verbose = True

        b = MajorDomoBroker(verbose)
        assert b.verbose == verbose
        assert b.services == {}
        assert b.workers == {}
        assert b.waiting == []
        assert b.heartbeat_at
        assert isinstance(b.ctx, zmq.Context)
        assert b.socket
        assert b.socket.linger == 0
        assert isinstance(b.poller, zmq.Poller)

    def test_broker_destroy(self, address):
        """Test destroy broker method"""
        b = MajorDomoBroker(False)
        b.require_worker(address)
        assert len(b.workers) == 1
        b.destroy()
        assert len(b.workers) == 0

    @pytest.mark.xfail
    #TODO building up to it
    def test_process_client(self):
        """Test process client method"""
        pass

    def test_require_worker(self, broker, broker_verbose, address):
        """Test require worker method"""
        assert len(broker.workers) == 0
        worker = broker.require_worker(address)
        assert len(broker.workers) == 1
        assert worker.address == address

        # verbose
        worker = broker_verbose.require_worker(address)
        assert len(broker_verbose.workers) == 1

    @pytest.mark.xfail
    #TODO building up to it
    def test_process_worker(self):
        """Test process worker method"""
        pass

    def test_delete_worker(self, broker, address, service):
        """Test delete worker method"""
        assert len(broker.workers) == 0
        worker = broker.require_worker(address)
        assert len(broker.workers) == 1

        # delete but no disconnect
        broker.delete_worker(worker, False)
        assert len(broker.workers) == 0

        # test when worker service is set
        worker = broker.require_worker(address)
        worker.service = broker.require_service(service)
        broker.worker_waiting(worker)
        assert len(broker.workers) == 1
        broker.delete_worker(worker, False)
        assert len(broker.workers) == 0

    def test_delete_worker_disconnect(self, broker, address):
        """Test delete worker method with disconnect call to worker"""
        assert len(broker.workers) == 0
        worker = broker.require_worker(address)
        assert len(broker.workers) == 1

        # delete and disconnect
        broker.delete_worker(worker, True)
        assert len(broker.workers) == 0

    def test_require_service(self, broker, address, service):
        """Test require service method"""
        assert len(broker.services) == 0
        worker = broker.require_worker(address)
        worker.service = broker.require_service(service)
        assert len(broker.services) == 1
        assert service in broker.services.keys()
        assert worker.service in broker.services.values()
        assert isinstance(broker.services[service], Service)

    def test_service_internal(self, broker):
        """Test service internal method"""
        broker.service_internal(b"mmi.service", [b"", b"Hello"])

    def test_send_heartbeats(self, broker, address, service):
        """Test send heartbeats method"""
        worker = broker.require_worker(address)
        worker.service = broker.require_service(service)
        broker.worker_waiting(worker)

        cur_heatbeat_at = broker.heartbeat_at
        broker.heartbeat_at = time.time() - 10
        broker.send_heartbeats()
        assert broker.heartbeat_at > cur_heatbeat_at

    def test_purge_workers(self, broker, address, service):
        """Test purge workers method"""
        worker1 = broker.require_worker(address)
        worker1.service = broker.require_service(service)

        worker2 = broker.require_worker(b"tcp://localhost:6667")
        worker2.service = broker.require_service(b"S_LOG")

        broker.worker_waiting(worker1)
        broker.worker_waiting(worker2)

        assert len(broker.waiting) == 2

        # modify expiry to be in past for worker1
        worker1.expiry = time.time() - 1

        broker.purge_workers()
        assert len(broker.waiting) == 1

    def test_worker_waiting(self, broker, address, service):
        """Test worker waiting method"""
        worker = broker.require_worker(address)
        worker.service = broker.require_service(service)
        assert len(broker.waiting) == 0
        assert len(worker.service.waiting) == 0
        broker.worker_waiting(worker)
        assert len(broker.waiting) == 1
        assert len(worker.service.waiting) == 1

    def test_dispatch(self, broker, address, service):
        """Test dispatch method"""
        worker = broker.require_worker(address)
        worker.service = broker.require_service(service)
        broker.worker_waiting(worker)
        assert len(worker.service.waiting) == 1
        broker.dispatch(worker.service, [b'hello'])
        assert len(worker.service.requests) == 0
        assert len(worker.service.waiting) == 0

    def test_send_to_worker(self, broker, broker_verbose, address):
        """Test send to worker method"""
        worker = broker.require_worker(address)
        assert broker.send_to_worker(worker, W_READY, None) is None
        assert broker.send_to_worker(worker, W_READY, None, b'Hello') is None
        assert broker.send_to_worker(worker, W_READY, b'Hello') is None

        worker = broker_verbose.require_worker(address)

        with support.captured_stdout() as s:
            assert broker_verbose.send_to_worker(worker, W_READY, None) is None
            assert "{addr}".format(addr=address) in s.getvalue()
<commit_msg>added process_client broker tests<commit_after>import logging
import pytest
import sys
import time
import zmq

from mock import Mock
from test import support

from mdbase.broker import Service, Worker, MajorDomoBroker, W_READY

log = logging.getLogger()
log.addHandler(logging.StreamHandler(sys.stdout))

@pytest.fixture
def broker():
    return MajorDomoBroker()

@pytest.fixture
def broker_verbose():
    return MajorDomoBroker(True)

@pytest.fixture
def address():
    return b"tcp://localhost:6666"

@pytest.fixture
def service():
    return Service(b"S_ECHO")

class TestBrokerService():
    def test_service_instantiate(self):
        """Test instantiating service model"""
        name = "Service1"
        s = Service(name)
        assert s.requests == []
        assert s.waiting == []
        assert s.name == name

class TestBrokerWorker():
    def test_worker_instantiate(self):
        """Test instantiating worker model"""
        identity = hex(1)
        address = "worker1"
        expiry = 3
        w = Worker(identity, address, expiry)
        assert w.service is None
        assert w.identity == identity
        assert w.address == address
        assert w.expiry is not None

class TestBrokerModel():
    def test_broker_instantiate(self):
        """Test instantiating majordomo worker model"""
        verbose = True

        b = MajorDomoBroker(verbose)
        assert b.verbose == verbose
        assert b.services == {}
        assert b.workers == {}
        assert b.waiting == []
        assert b.heartbeat_at
        assert isinstance(b.ctx, zmq.Context)
        assert b.socket
        assert b.socket.linger == 0
        assert isinstance(b.poller, zmq.Poller)

    def test_broker_destroy(self, address):
        """Test destroy broker method"""
        b = MajorDomoBroker(False)
        b.require_worker(address)
        assert len(b.workers) == 1
        b.destroy()
        assert len(b.workers) == 0

    def test_process_client(self, broker):
        """Test process client method with service internal"""
        service_internal_mock = Mock()
        broker.service_internal = service_internal_mock
        broker.process_client("TEST", [broker.INTERNAL_SERVICE_PREFIX, "hello"])
        service_internal_mock.assert_called_with(broker.INTERNAL_SERVICE_PREFIX, ["TEST", "", "hello"])

    def test_process_client_dispatch(self, broker):
        """Test process client method with service internal"""
        dispatch_mock = Mock()
        broker.dispatch = dispatch_mock
        broker.process_client("TEST", ["srv1", "hello"])
        dispatch_mock.assert_called_with(broker.require_service("srv1"), ["TEST", "", "hello"])

    def test_require_worker(self, broker, broker_verbose, address):
        """Test require worker method"""
        assert len(broker.workers) == 0
        worker = broker.require_worker(address)
        assert len(broker.workers) == 1
        assert worker.address == address

        # verbose
        worker = broker_verbose.require_worker(address)
        assert len(broker_verbose.workers) == 1

    @pytest.mark.xfail
    #TODO building up to it
    def test_process_worker(self):
        """Test process worker method"""
        pass

    def test_delete_worker(self, broker, address, service):
        """Test delete worker method"""
        assert len(broker.workers) == 0
        worker = broker.require_worker(address)
        assert len(broker.workers) == 1

        # delete but no disconnect
        broker.delete_worker(worker, False)
        assert len(broker.workers) == 0

        # test when worker service is set
        worker = broker.require_worker(address)
        worker.service = broker.require_service(service)
        broker.worker_waiting(worker)
        assert len(broker.workers) == 1
        broker.delete_worker(worker, False)
        assert len(broker.workers) == 0

    def test_delete_worker_disconnect(self, broker, address):
        """Test delete worker method with disconnect call to worker"""
        assert len(broker.workers) == 0
        worker = broker.require_worker(address)
        assert len(broker.workers) == 1

        # delete and disconnect
        broker.delete_worker(worker, True)
        assert len(broker.workers) == 0

    def test_require_service(self, broker, address, service):
        """Test require service method"""
        assert len(broker.services) == 0
        worker = broker.require_worker(address)
        worker.service = broker.require_service(service)
        assert len(broker.services) == 1
        assert service in broker.services.keys()
        assert worker.service in broker.services.values()
        assert isinstance(broker.services[service], Service)

    def test_service_internal(self, broker):
        """Test service internal method"""
        broker.service_internal(b"mmi.service", [b"", b"Hello"])

    def test_send_heartbeats(self, broker, address, service):
        """Test send heartbeats method"""
        worker = broker.require_worker(address)
        worker.service = broker.require_service(service)
        broker.worker_waiting(worker)

        cur_heatbeat_at = broker.heartbeat_at
        broker.heartbeat_at = time.time() - 10
        broker.send_heartbeats()
        assert broker.heartbeat_at > cur_heatbeat_at

    def test_purge_workers(self, broker, address, service):
        """Test purge workers method"""
        worker1 = broker.require_worker(address)
        worker1.service = broker.require_service(service)

        worker2 = broker.require_worker(b"tcp://localhost:6667")
        worker2.service = broker.require_service(b"S_LOG")

        broker.worker_waiting(worker1)
        broker.worker_waiting(worker2)

        assert len(broker.waiting) == 2

        # modify expiry to be in past for worker1
        worker1.expiry = time.time() - 1

        broker.purge_workers()
        assert len(broker.waiting) == 1

    def test_worker_waiting(self, broker, address, service):
        """Test worker waiting method"""
        worker = broker.require_worker(address)
        worker.service = broker.require_service(service)
        assert len(broker.waiting) == 0
        assert len(worker.service.waiting) == 0
        broker.worker_waiting(worker)
        assert len(broker.waiting) == 1
        assert len(worker.service.waiting) == 1

    def test_dispatch(self, broker, address, service):
        """Test dispatch method"""
        worker = broker.require_worker(address)
        worker.service = broker.require_service(service)
        broker.worker_waiting(worker)
        assert len(worker.service.waiting) == 1
        broker.dispatch(worker.service, [b'hello'])
        assert len(worker.service.requests) == 0
        assert len(worker.service.waiting) == 0

    def test_send_to_worker(self, broker, broker_verbose, address):
        """Test send to worker method"""
        worker = broker.require_worker(address)
        assert broker.send_to_worker(worker, W_READY, None) is None
        assert broker.send_to_worker(worker, W_READY, None, b'Hello') is None
        assert broker.send_to_worker(worker, W_READY, b'Hello') is None

        worker = broker_verbose.require_worker(address)

        with support.captured_stdout() as s:
            assert broker_verbose.send_to_worker(worker, W_READY, None) is None
            assert "{addr}".format(addr=address) in s.getvalue()
<|endoftext|>


<commit_before>"""Unit tests for mbedtls.cipher."""

# Disable checks for violations that are acceptable in tests.
# pylint: disable=missing-docstring
# pylint: disable=attribute-defined-outside-init
# pylint: disable=invalid-name, redefined-outer-name

import pickle
import sys
from collections import defaultdict
from typing import DefaultDict, Mapping, NamedTuple, Sequence

import pytest  # type: ignore

from mbedtls.cipher import (  # type: ignore
    AES,
    ARC4,
    ARIA,
    CHACHA20,
    DES,
    DES3,
    MODE_CBC,
    MODE_CFB,
    Blowfish,
    Camellia,
    Cipher,
    DES3dbl,
    Mode,
    get_supported_ciphers,
)
from mbedtls.cipher._cipher import CIPHER_NAME  # type: ignore
from mbedtls.exceptions import TLSError  # type: ignore

class Size(NamedTuple):
    key_size: Sequence[int]
    iv_size: int

def constant(value):
    return lambda: value

SUPPORTED_SIZES: Mapping[str, DefaultDict[Mode, Size]] = {
    "mbedtls.cipher.AES": defaultdict(
        constant(Size(key_size=(16, 24, 32), iv_size=16)),
        {
            Mode.CCM: Size((16, 24, 32), 12),
            Mode.ECB: Size((16, 24, 32), 0),
            Mode.GCM: Size((16, 24, 32), 12),
            Mode.XTS: Size((32, 64), 16),
        },
    ),
    "mbedtls.cipher.ARC4": defaultdict(
        constant(Size(key_size=(ARC4.key_size,), iv_size=0)),
    ),
    "mbedtls.cipher.ARIA": defaultdict(
        constant(Size(key_size=(16, 24, 32), iv_size=16)),
        {
            Mode.CBC: Size((16, 24, 32), 16),
            Mode.ECB: Size((16, 24, 32), 0),
            Mode.GCM: Size((16, 24, 32), 12),
        },
    ),
    "mbedtls.cipher.Blowfish": defaultdict(
        constant(Size(key_size=(16,), iv_size=8)),
        {Mode.ECB: Size((16,), 0)},
    ),
    "mbedtls.cipher.Camellia": defaultdict(
        constant(Size(key_size=(16, 24, 32), iv_size=16)),
        {
            Mode.ECB: Size((16, 24, 32), 0),
            Mode.GCM: Size((16, 24, 32), 12),
        },
    ),
    "mbedtls.cipher.CHACHA20": defaultdict(
        constant(Size(key_size=(CHACHA20.key_size,), iv_size=12)),
        {Mode.ECB: Size((CHACHA20.key_size,), 0)},
    ),
    "mbedtls.cipher.DES": defaultdict(
        constant(Size(key_size=(DES.key_size,), iv_size=8)),
        {Mode.ECB: Size((DES.key_size,), 0)},
    ),
    "mbedtls.cipher.DES3": defaultdict(
        constant(Size(key_size=(DES3.key_size,), iv_size=8)),
        {Mode.ECB: Size((DES3.key_size,), 0)},
    ),
    "mbedtls.cipher.DES3dbl": defaultdict(
        constant(Size(key_size=(DES3dbl.key_size,), iv_size=8)),
        {Mode.ECB: Size((DES3dbl.key_size,), 0)},
    ),
}

SUPPORTED_MODES: Mapping[str, Sequence[Mode]] = {
    "mbedtls.cipher.AES": (
        Mode.ECB,
        Mode.CBC,
        Mode.CFB,
        Mode.CTR,
        Mode.OFB,
        Mode.XTS,
    ),
    "mbedtls.cipher.ARC4": (Mode.STREAM,),
    "mbedtls.cipher.ARIA": (Mode.ECB, Mode.CBC, Mode.CTR, Mode.GCM),
    "mbedtls.cipher.Blowfish": (Mode.ECB, Mode.CBC, Mode.CFB, Mode.CTR),
    "mbedtls.cipher.Camellia": (
        Mode.ECB,
        Mode.CBC,
        Mode.CFB,
        Mode.CTR,
        Mode.GCM,
    ),
    "mbedtls.cipher.CHACHA20": (Mode.STREAM,),
    "mbedtls.cipher.DES": (Mode.ECB, Mode.CBC),
    "mbedtls.cipher.DES3": (Mode.ECB, Mode.CBC),
    "mbedtls.cipher.DES3dbl": (Mode.ECB, Mode.CBC),
}

SUPPORTED_AEAD_MODES: Mapping[str, Sequence[Mode]] = {
    "mbedtls.cipher.AES": (Mode.GCM, Mode.CCM),
    "mbedtls.cipher.CHACHA20": (Mode.CHACHAPOLY,),
}

def gen_cipher_data(module, *, modes: Mapping[str, Sequence[Mode]]):
    for mode in modes[module.__name__]:
        sizes = SUPPORTED_SIZES[module.__name__][mode]
        for key_size in sizes.key_size:
            yield key_size, mode, sizes.iv_size

def test_cipher_list():
    assert len(CIPHER_NAME) == 74

def test_get_supported_ciphers():
    cl = get_supported_ciphers()
    assert cl and set(cl).issubset(set(CIPHER_NAME))

def test_wrong_size_raises_exception():
    with pytest.raises(NotImplementedError):
        Cipher(b"AES-512-ECB", b"", Mode.ECB, b"")

def test_random_name_raises_exception():
    with pytest.raises(NotImplementedError):
        Cipher(b"RANDOM TEXT IS NOT A CIPHER", b"", Mode.ECB, b"")

def test_zero_length_raises_exception():
    with pytest.raises(NotImplementedError):
        Cipher(b"", b"", Mode.ECB, b"")

@pytest.mark.parametrize("mode", [MODE_CBC, Mode.CBC])
def test_cbc_raises_value_error_without_iv(mode):
    with pytest.raises(ValueError):
        Cipher(b"AES-512-CBC", b"", mode, b"")

@pytest.mark.parametrize("mode", [MODE_CFB, Mode.CFB])
def test_cfb_raises_value_error_without_iv(mode):
    with pytest.raises(ValueError):
        Cipher(b"AES-512-CFB", b"", mode, b"")

class TestCipher:
    @pytest.fixture(
        params=[
            AES,
            ARC4,
            ARIA,
            Blowfish,
            Camellia,
            CHACHA20,
            DES,
            DES3,
            DES3dbl,
        ]
    )
    def module(self, request):
        if request.param is ARIA and sys.platform.startswith("win"):
            return pytest.skip()
        return request.param

    def test_pickle(self, module, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_MODES
        ):
            cipher = module.new(randbytes(key_size), mode, randbytes(iv_size))
            with pytest.raises(TypeError) as excinfo:
                pickle.dumps(cipher)

            assert str(excinfo.value).startswith("cannot pickle")

    def test_accessors(self, module, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_MODES
        ):
            cipher = module.new(randbytes(key_size), mode, randbytes(iv_size))
            assert cipher.key_size == key_size
            assert cipher.mode == mode
            assert cipher.iv_size == iv_size
            assert module.block_size == cipher.block_size
            assert module.key_size in {module.key_size, None}

    def test_cipher_name(self, module, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_MODES
        ):
            cipher = module.new(randbytes(key_size), mode, randbytes(iv_size))
            assert cipher.name in CIPHER_NAME
            assert CIPHER_NAME[cipher._type] == cipher.name
            assert str(cipher) == cipher.name.decode("ascii")

    def test_encrypt_decrypt(self, module, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_MODES
        ):
            cipher = module.new(randbytes(key_size), mode, randbytes(iv_size))
            data = randbytes(
                cipher.block_size
                if cipher.mode is Mode.ECB
                else cipher.block_size * 128
            )
            assert cipher.decrypt(cipher.encrypt(data)) == data

    def test_encrypt_nothing_raises(self, module, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_MODES
        ):
            cipher = module.new(randbytes(key_size), mode, randbytes(iv_size))
            with pytest.raises(TLSError):
                cipher.encrypt(b"")

    def test_decrypt_nothing_raises(self, module, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_MODES
        ):
            cipher = module.new(randbytes(key_size), mode, randbytes(iv_size))
            with pytest.raises(TLSError):
                cipher.decrypt(b"")

    def test_cbc_requires_padding(self, module, randbytes):
        mode = Mode.CBC
        if mode not in SUPPORTED_MODES[module.__name__]:
            return pytest.skip(
                f"unsupported mode for {module.__name__!r}: {mode!s}"
            )

        sizes = SUPPORTED_SIZES[module.__name__][mode]
        for key_size in sizes.key_size:
            cipher = module.new(
                randbytes(key_size), mode, iv=randbytes(sizes.iv_size)
            )
            data = randbytes(
                cipher.block_size
                if cipher.mode is Mode.ECB
                else cipher.block_size * 128
            )
            data += b"\0"
            if cipher.mode is Mode.CBC:
                with pytest.raises(ValueError):
                    cipher.encrypt(data)

class TestAEADCipher:
    @pytest.fixture(params=[AES, CHACHA20])
    def module(self, request):
        return request.param

    @pytest.mark.parametrize("ad_size", [0, 1, 16, 256])
    def test_pickle(self, module, ad_size, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_AEAD_MODES
        ):
            cipher = module.new(
                randbytes(key_size),
                mode,
                iv=randbytes(iv_size),
                ad=randbytes(ad_size),
            )
            with pytest.raises(TypeError) as excinfo:
                pickle.dumps(cipher)

            assert str(excinfo.value).startswith("cannot pickle")

    @pytest.mark.parametrize("ad_size", [0, 1, 16, 256])
    def test_accessors(self, module, ad_size, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_AEAD_MODES
        ):
            cipher = module.new(
                randbytes(key_size),
                mode,
                iv=randbytes(iv_size),
                ad=randbytes(ad_size),
            )
            assert cipher.key_size == key_size
            assert cipher.mode == mode
            assert cipher.iv_size == iv_size
            assert module.block_size == cipher.block_size
            assert module.key_size in {module.key_size, None}

    @pytest.mark.parametrize("ad_size", [0, 1, 16, 256])
    def test_cipher_name(self, module, ad_size, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_AEAD_MODES
        ):
            cipher = module.new(
                randbytes(key_size),
                mode,
                iv=randbytes(iv_size),
                ad=randbytes(ad_size),
            )
            assert cipher.name in CIPHER_NAME
            assert CIPHER_NAME[cipher._type] == cipher.name
            assert str(cipher) == cipher.name.decode("ascii")

    @pytest.mark.parametrize("ad_size", [0, 1, 16, 256])
    def test_encrypt_decrypt(self, module, ad_size, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_AEAD_MODES
        ):
            cipher = module.new(
                randbytes(key_size),
                mode,
                iv=randbytes(iv_size),
                ad=randbytes(ad_size),
            )
            data = randbytes(
                cipher.block_size
                if cipher.mode is Mode.ECB
                else cipher.block_size * 128
            )
            msg, tag = cipher.encrypt(data)
            assert cipher.decrypt(msg, tag) == data

    @pytest.mark.parametrize("ad_size", [0, 1, 16, 256])
    def test_encrypt_nothing_raises(self, module, ad_size, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_AEAD_MODES
        ):
            cipher = module.new(
                randbytes(key_size),
                mode,
                iv=randbytes(iv_size),
                ad=randbytes(ad_size),
            )
            with pytest.raises(TLSError):
                cipher.encrypt(b"")

    @pytest.mark.parametrize("ad_size", [0, 1, 16, 256])
    def test_decrypt_nothing_raises(self, module, ad_size, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_AEAD_MODES
        ):
            cipher = module.new(
                randbytes(key_size),
                mode,
                iv=randbytes(iv_size),
                ad=randbytes(ad_size),
            )
            data = randbytes(
                cipher.block_size
                if cipher.mode is Mode.ECB
                else cipher.block_size * 128
            )
            msg, tag = cipher.encrypt(data)
            with pytest.raises(TLSError):
                cipher.decrypt(b"", tag)
<commit_msg>test_cipher: Improve typing<commit_after>"""Unit tests for mbedtls.cipher."""

# Disable checks for violations that are acceptable in tests.
# pylint: disable=missing-docstring
# pylint: disable=attribute-defined-outside-init
# pylint: disable=invalid-name, redefined-outer-name

import pickle
import sys
from collections import defaultdict
from typing import (
    Callable,
    Iterator,
    Mapping,
    NamedTuple,
    Sequence,
    Tuple,
    TypeVar,
)

import pytest  # type: ignore

from mbedtls.cipher import (  # type: ignore
    AES,
    ARC4,
    ARIA,
    CHACHA20,
    DES,
    DES3,
    MODE_CBC,
    MODE_CFB,
    Blowfish,
    Camellia,
    Cipher,
    DES3dbl,
    Mode,
    get_supported_ciphers,
)
from mbedtls.cipher._cipher import CIPHER_NAME  # type: ignore
from mbedtls.exceptions import TLSError  # type: ignore

class Size(NamedTuple):
    key_size: Sequence[int]
    iv_size: int

T = TypeVar("T")

def constant(value: T) -> Callable[[], T]:
    return lambda: value

SUPPORTED_SIZES: Mapping[str, Mapping[Mode, Size]] = {
    "mbedtls.cipher.AES": defaultdict(
        constant(Size(key_size=(16, 24, 32), iv_size=16)),
        {
            Mode.CCM: Size((16, 24, 32), 12),
            Mode.ECB: Size((16, 24, 32), 0),
            Mode.GCM: Size((16, 24, 32), 12),
            Mode.XTS: Size((32, 64), 16),
        },
    ),
    "mbedtls.cipher.ARC4": defaultdict(
        constant(Size(key_size=(ARC4.key_size,), iv_size=0)),
    ),
    "mbedtls.cipher.ARIA": defaultdict(
        constant(Size(key_size=(16, 24, 32), iv_size=16)),
        {
            Mode.CBC: Size((16, 24, 32), 16),
            Mode.ECB: Size((16, 24, 32), 0),
            Mode.GCM: Size((16, 24, 32), 12),
        },
    ),
    "mbedtls.cipher.Blowfish": defaultdict(
        constant(Size(key_size=(16,), iv_size=8)),
        {Mode.ECB: Size((16,), 0)},
    ),
    "mbedtls.cipher.Camellia": defaultdict(
        constant(Size(key_size=(16, 24, 32), iv_size=16)),
        {
            Mode.ECB: Size((16, 24, 32), 0),
            Mode.GCM: Size((16, 24, 32), 12),
        },
    ),
    "mbedtls.cipher.CHACHA20": defaultdict(
        constant(Size(key_size=(CHACHA20.key_size,), iv_size=12)),
        {Mode.ECB: Size((CHACHA20.key_size,), 0)},
    ),
    "mbedtls.cipher.DES": defaultdict(
        constant(Size(key_size=(DES.key_size,), iv_size=8)),
        {Mode.ECB: Size((DES.key_size,), 0)},
    ),
    "mbedtls.cipher.DES3": defaultdict(
        constant(Size(key_size=(DES3.key_size,), iv_size=8)),
        {Mode.ECB: Size((DES3.key_size,), 0)},
    ),
    "mbedtls.cipher.DES3dbl": defaultdict(
        constant(Size(key_size=(DES3dbl.key_size,), iv_size=8)),
        {Mode.ECB: Size((DES3dbl.key_size,), 0)},
    ),
}

SUPPORTED_MODES: Mapping[str, Sequence[Mode]] = {
    "mbedtls.cipher.AES": (
        Mode.ECB,
        Mode.CBC,
        Mode.CFB,
        Mode.CTR,
        Mode.OFB,
        Mode.XTS,
    ),
    "mbedtls.cipher.ARC4": (Mode.STREAM,),
    "mbedtls.cipher.ARIA": (Mode.ECB, Mode.CBC, Mode.CTR, Mode.GCM),
    "mbedtls.cipher.Blowfish": (Mode.ECB, Mode.CBC, Mode.CFB, Mode.CTR),
    "mbedtls.cipher.Camellia": (
        Mode.ECB,
        Mode.CBC,
        Mode.CFB,
        Mode.CTR,
        Mode.GCM,
    ),
    "mbedtls.cipher.CHACHA20": (Mode.STREAM,),
    "mbedtls.cipher.DES": (Mode.ECB, Mode.CBC),
    "mbedtls.cipher.DES3": (Mode.ECB, Mode.CBC),
    "mbedtls.cipher.DES3dbl": (Mode.ECB, Mode.CBC),
}

SUPPORTED_AEAD_MODES: Mapping[str, Sequence[Mode]] = {
    "mbedtls.cipher.AES": (Mode.GCM, Mode.CCM),
    "mbedtls.cipher.CHACHA20": (Mode.CHACHAPOLY,),
}

def gen_cipher_data(
    module: type, *, modes: Mapping[str, Sequence[Mode]]
) -> Iterator[Tuple[int, Mode, int]]:
    for mode in modes[module.__name__]:
        sizes = SUPPORTED_SIZES[module.__name__][mode]
        for key_size in sizes.key_size:
            yield key_size, mode, sizes.iv_size

def test_cipher_list() -> None:
    assert len(CIPHER_NAME) == 74

def test_get_supported_ciphers() -> None:
    cl = get_supported_ciphers()
    assert cl and set(cl).issubset(set(CIPHER_NAME))

def test_wrong_size_raises_exception() -> None:
    with pytest.raises(NotImplementedError):
        Cipher(b"AES-512-ECB", b"", Mode.ECB, b"")

def test_random_name_raises_exception() -> None:
    with pytest.raises(NotImplementedError):
        Cipher(b"RANDOM TEXT IS NOT A CIPHER", b"", Mode.ECB, b"")

def test_zero_length_raises_exception() -> None:
    with pytest.raises(NotImplementedError):
        Cipher(b"", b"", Mode.ECB, b"")

@pytest.mark.parametrize("mode", [MODE_CBC, Mode.CBC])
def test_cbc_raises_value_error_without_iv(mode: Mode) -> None:
    with pytest.raises(ValueError):
        Cipher(b"AES-512-CBC", b"", mode, b"")

@pytest.mark.parametrize("mode", [MODE_CFB, Mode.CFB])
def test_cfb_raises_value_error_without_iv(mode: Mode) -> None:
    with pytest.raises(ValueError):
        Cipher(b"AES-512-CFB", b"", mode, b"")

class TestCipher:
    @pytest.fixture(
        params=[
            AES,
            ARC4,
            ARIA,
            Blowfish,
            Camellia,
            CHACHA20,
            DES,
            DES3,
            DES3dbl,
        ]
    )
    def module(self, request):
        if request.param is ARIA and sys.platform.startswith("win"):
            return pytest.skip()
        return request.param

    def test_pickle(self, module, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_MODES
        ):
            cipher = module.new(randbytes(key_size), mode, randbytes(iv_size))
            with pytest.raises(TypeError) as excinfo:
                pickle.dumps(cipher)

            assert str(excinfo.value).startswith("cannot pickle")

    def test_accessors(self, module, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_MODES
        ):
            cipher = module.new(randbytes(key_size), mode, randbytes(iv_size))
            assert cipher.key_size == key_size
            assert cipher.mode == mode
            assert cipher.iv_size == iv_size
            assert module.block_size == cipher.block_size
            assert module.key_size in {module.key_size, None}

    def test_cipher_name(self, module, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_MODES
        ):
            cipher = module.new(randbytes(key_size), mode, randbytes(iv_size))
            assert cipher.name in CIPHER_NAME
            assert CIPHER_NAME[cipher._type] == cipher.name
            assert str(cipher) == cipher.name.decode("ascii")

    def test_encrypt_decrypt(self, module, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_MODES
        ):
            cipher = module.new(randbytes(key_size), mode, randbytes(iv_size))
            data = randbytes(
                cipher.block_size
                if cipher.mode is Mode.ECB
                else cipher.block_size * 128
            )
            assert cipher.decrypt(cipher.encrypt(data)) == data

    def test_encrypt_nothing_raises(self, module, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_MODES
        ):
            cipher = module.new(randbytes(key_size), mode, randbytes(iv_size))
            with pytest.raises(TLSError):
                cipher.encrypt(b"")

    def test_decrypt_nothing_raises(self, module, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_MODES
        ):
            cipher = module.new(randbytes(key_size), mode, randbytes(iv_size))
            with pytest.raises(TLSError):
                cipher.decrypt(b"")

    def test_cbc_requires_padding(self, module, randbytes):
        mode = Mode.CBC
        if mode not in SUPPORTED_MODES[module.__name__]:
            return pytest.skip(
                f"unsupported mode for {module.__name__!r}: {mode!s}"
            )

        sizes = SUPPORTED_SIZES[module.__name__][mode]
        for key_size in sizes.key_size:
            cipher = module.new(
                randbytes(key_size), mode, iv=randbytes(sizes.iv_size)
            )
            data = randbytes(
                cipher.block_size
                if cipher.mode is Mode.ECB
                else cipher.block_size * 128
            )
            data += b"\0"
            if cipher.mode is Mode.CBC:
                with pytest.raises(ValueError):
                    cipher.encrypt(data)

class TestAEADCipher:
    @pytest.fixture(params=[AES, CHACHA20])
    def module(self, request):
        return request.param

    @pytest.mark.parametrize("ad_size", [0, 1, 16, 256])
    def test_pickle(self, module, ad_size, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_AEAD_MODES
        ):
            cipher = module.new(
                randbytes(key_size),
                mode,
                iv=randbytes(iv_size),
                ad=randbytes(ad_size),
            )
            with pytest.raises(TypeError) as excinfo:
                pickle.dumps(cipher)

            assert str(excinfo.value).startswith("cannot pickle")

    @pytest.mark.parametrize("ad_size", [0, 1, 16, 256])
    def test_accessors(self, module, ad_size, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_AEAD_MODES
        ):
            cipher = module.new(
                randbytes(key_size),
                mode,
                iv=randbytes(iv_size),
                ad=randbytes(ad_size),
            )
            assert cipher.key_size == key_size
            assert cipher.mode == mode
            assert cipher.iv_size == iv_size
            assert module.block_size == cipher.block_size
            assert module.key_size in {module.key_size, None}

    @pytest.mark.parametrize("ad_size", [0, 1, 16, 256])
    def test_cipher_name(self, module, ad_size, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_AEAD_MODES
        ):
            cipher = module.new(
                randbytes(key_size),
                mode,
                iv=randbytes(iv_size),
                ad=randbytes(ad_size),
            )
            assert cipher.name in CIPHER_NAME
            assert CIPHER_NAME[cipher._type] == cipher.name
            assert str(cipher) == cipher.name.decode("ascii")

    @pytest.mark.parametrize("ad_size", [0, 1, 16, 256])
    def test_encrypt_decrypt(self, module, ad_size, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_AEAD_MODES
        ):
            cipher = module.new(
                randbytes(key_size),
                mode,
                iv=randbytes(iv_size),
                ad=randbytes(ad_size),
            )
            data = randbytes(
                cipher.block_size
                if cipher.mode is Mode.ECB
                else cipher.block_size * 128
            )
            msg, tag = cipher.encrypt(data)
            assert cipher.decrypt(msg, tag) == data

    @pytest.mark.parametrize("ad_size", [0, 1, 16, 256])
    def test_encrypt_nothing_raises(self, module, ad_size, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_AEAD_MODES
        ):
            cipher = module.new(
                randbytes(key_size),
                mode,
                iv=randbytes(iv_size),
                ad=randbytes(ad_size),
            )
            with pytest.raises(TLSError):
                cipher.encrypt(b"")

    @pytest.mark.parametrize("ad_size", [0, 1, 16, 256])
    def test_decrypt_nothing_raises(self, module, ad_size, randbytes):
        for key_size, mode, iv_size in gen_cipher_data(
            module, modes=SUPPORTED_AEAD_MODES
        ):
            cipher = module.new(
                randbytes(key_size),
                mode,
                iv=randbytes(iv_size),
                ad=randbytes(ad_size),
            )
            data = randbytes(
                cipher.block_size
                if cipher.mode is Mode.ECB
                else cipher.block_size * 128
            )
            msg, tag = cipher.encrypt(data)
            with pytest.raises(TLSError):
                cipher.decrypt(b"", tag)
<|endoftext|>


<commit_before>
from __future__ import division

import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(__file__)))

import time
from subprocess import Popen, PIPE

from tomcrypt.cipher import *
Descriptor = CipherDescriptor

internal_tests = (
    ( 16,
      ( 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
        0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f ),
      ( 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77,
        0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff ),
      ( 0x69, 0xc4, 0xe0, 0xd8, 0x6a, 0x7b, 0x04, 0x30,
        0xd8, 0xcd, 0xb7, 0x80, 0x70, 0xb4, 0xc5, 0x5a )
    ), (
      24,
      ( 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
        0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
        0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17 ),
      ( 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77,
        0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff ),
      ( 0xdd, 0xa9, 0x7c, 0xa4, 0x86, 0x4c, 0xdf, 0xe0,
        0x6e, 0xaf, 0x70, 0xa0, 0xec, 0x0d, 0x71, 0x91 )
    ), (
      32,
      ( 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
        0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
        0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
        0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f ),
      ( 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77,
        0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff ),
      ( 0x8e, 0xa2, 0xb7, 0xca, 0x51, 0x67, 0x45, 0xbf,
        0xea, 0xfc, 0x49, 0x90, 0x4b, 0x49, 0x60, 0x89 )
    )
)

def test_tomcrypt():
    test_cipher()

def test_internal():
	for keysize, key, pt, ct in internal_tests:
		key = ''.join(map(chr, key))
		pt = ''.join(map(chr, pt))
		ct = ''.join(map(chr, ct))

		cipher = Cipher(key=key, cipher='aes', mode='ecb')
		test_ct = cipher.encrypt(pt)
		assert ct == test_ct, 'internal encrypt: %s != %s' % (ct.encode('hex'), test_ct.encode('hex'))
		cipher = Cipher(key=key, cipher='aes', mode='ecb')
		test_pt = cipher.decrypt(ct)
		assert pt == test_pt, 'internal decrypt: %s != %s' % (pt.encode('hex'), test_pt.encode('hex'))

def test_speed():
	cipher = Cipher('0123456789abcdef', cipher='aes', mode='ecb')
	start_time = time.time()
	txt = '0123456789abcdef'
	for i in xrange(50000):
		txt = cipher.encrypt(txt)
	for i in xrange(50000):
		txt = cipher.decrypt(txt)
	print 'Each AES block done in %.2fns' % ((time.time() - start_time) * 10**9 / 10**5)
	assert txt == '0123456789abcdef', 'speed test is wrong: %r' % txt

def test_openssl():
	for cipher_name in 'aes', 'des':
		cipher_desc = Descriptor(cipher=cipher_name)
		keysizes = []
		for i in range(cipher_desc.min_key_size, cipher_desc.max_key_size + 1):
			keysizes.append(cipher_desc.key_size(i))
		keysizes = list(sorted(set(keysizes)))
		for mode in 'ecb', 'cbc', 'cfb', 'ofb':
			for keysize in keysizes:
				keysize *= 8
				print cipher_name, keysize, mode
				for i in xrange(0, 10, 4):
					key = os.urandom(keysize//8)
					iv  = os.urandom(cipher_desc.block_size)
					pt  = os.urandom(i * 128 // 8)
					if cipher_name == 'aes':
						cipher_spec = 'aes-%d-%s' % (keysize, mode)
					elif cipher_name == 'des':
						cipher_spec = 'des-%s' % mode
					elif cipher_name == 'blowfish':
						cipher_spec = 'bf-%s' % mode
					proc = Popen(('openssl enc -e -%s -nopad -nosalt -K %s -iv %s' % (cipher_spec, key.encode('hex'), iv.encode('hex'))).split(), stdin=PIPE, stdout=PIPE, stderr=PIPE)
					out, err = proc.communicate(pt)
					assert not err, err
					cipher = Cipher(key=key, iv=iv, cipher=cipher_name, mode=mode)
					ct = cipher.encrypt(pt)
					assert ct == out, 'openssl: %s != %s' % (ct.encode('hex'), out.encode('hex'))

def test_external():
	for filename in os.listdir('test_vectors'):
		if 'CFB1' in filename:
			continue
		mode = filename[:3].lower()
		fh = open('test_vectors/' + filename, 'rb')
		type = fh.readline().strip()[1:-1].lower()
		fh.readline()

		# print mode, type, filename
		data = {}
		for line in fh:
			line = line.strip()
			if not line and data:
				# print mode, type, data['count']
				# for k in 'iv', 'key', 'plaintext', 'ciphertext':
				# 	if k in data:
				# 		print '\t%3s: %s' % (k, data.get(k))
				key = data['key'].decode('hex')
				iv  = data.get('iv', '').decode('hex')
				pt  = data['plaintext'].decode('hex')
				ct  = data['ciphertext'].decode('hex')
				cipher = Cipher(key=key, iv=iv or None, cipher='aes', mode=mode)
				if type == 'encrypt':
					res = cipher.encrypt(pt)
					assert res == ct, '%s #%s: %s != %s' % (name, data['count'], res.encode('hex'), data['ciphertext'])
				data = {}
			if ' = ' not in line:
				continue
			k, v = line.lower().split(' = ')
			data[k] = v

if __name__ == '__main__':
	start_time = time.time()
	print 'Running internal tests...'
	test_tomcrypt()
	print 'Running cython tests...'
	test_internal()
	print 'Running external tests...'
	test_external()
	print 'Running against OpenSSL...'
	test_openssl()
	print 'Running speed test...'
	test_speed()
	print 'Ran all tests in %.2fms' % (1000 * (time.time() - start_time))

<commit_msg>Remove one of the cipher tests.<commit_after>
from __future__ import division

import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(__file__)))

import time
from subprocess import Popen, PIPE

from tomcrypt.cipher import *
Descriptor = CipherDescriptor

def test_tomcrypt():
    test_cipher()

def test_speed():
	cipher = Cipher('0123456789abcdef', cipher='aes', mode='ecb')
	start_time = time.time()
	txt = '0123456789abcdef'
	for i in xrange(50000):
		txt = cipher.encrypt(txt)
	for i in xrange(50000):
		txt = cipher.decrypt(txt)
	print 'Each AES block done in %.2fns' % ((time.time() - start_time) * 10**9 / 10**5)
	assert txt == '0123456789abcdef', 'speed test is wrong: %r' % txt

def test_openssl():
	for cipher_name in 'aes', 'des':
		cipher_desc = Descriptor(cipher=cipher_name)
		keysizes = []
		for i in range(cipher_desc.min_key_size, cipher_desc.max_key_size + 1):
			keysizes.append(cipher_desc.key_size(i))
		keysizes = list(sorted(set(keysizes)))
		for mode in 'ecb', 'cbc', 'cfb', 'ofb':
			for keysize in keysizes:
				keysize *= 8
				print cipher_name, keysize, mode
				for i in xrange(0, 10, 4):
					key = os.urandom(keysize//8)
					iv  = os.urandom(cipher_desc.block_size)
					pt  = os.urandom(i * 128 // 8)
					if cipher_name == 'aes':
						cipher_spec = 'aes-%d-%s' % (keysize, mode)
					elif cipher_name == 'des':
						cipher_spec = 'des-%s' % mode
					elif cipher_name == 'blowfish':
						cipher_spec = 'bf-%s' % mode
					proc = Popen(('openssl enc -e -%s -nopad -nosalt -K %s -iv %s' % (cipher_spec, key.encode('hex'), iv.encode('hex'))).split(), stdin=PIPE, stdout=PIPE, stderr=PIPE)
					out, err = proc.communicate(pt)
					assert not err, err
					cipher = Cipher(key=key, iv=iv, cipher=cipher_name, mode=mode)
					ct = cipher.encrypt(pt)
					assert ct == out, 'openssl: %s != %s' % (ct.encode('hex'), out.encode('hex'))

def test_external():
	for filename in os.listdir('test_vectors'):
		if 'CFB1' in filename:
			continue
		mode = filename[:3].lower()
		fh = open('test_vectors/' + filename, 'rb')
		type = fh.readline().strip()[1:-1].lower()
		fh.readline()

		# print mode, type, filename
		data = {}
		for line in fh:
			line = line.strip()
			if not line and data:
				# print mode, type, data['count']
				# for k in 'iv', 'key', 'plaintext', 'ciphertext':
				# 	if k in data:
				# 		print '\t%3s: %s' % (k, data.get(k))
				key = data['key'].decode('hex')
				iv  = data.get('iv', '').decode('hex')
				pt  = data['plaintext'].decode('hex')
				ct  = data['ciphertext'].decode('hex')
				cipher = Cipher(key=key, iv=iv or None, cipher='aes', mode=mode)
				if type == 'encrypt':
					res = cipher.encrypt(pt)
					assert res == ct, '%s #%s: %s != %s' % (name, data['count'], res.encode('hex'), data['ciphertext'])
				data = {}
			if ' = ' not in line:
				continue
			k, v = line.lower().split(' = ')
			data[k] = v

if __name__ == '__main__':
	start_time = time.time()
	print 'Running tomcrypt tests...'
	test_tomcrypt()
	print 'Running known tests...'
	test_external()
	print 'Running against OpenSSL...'
	test_openssl()
	print 'Running speed test...'
	test_speed()
	print 'Ran all tests in %.2fms' % (1000 * (time.time() - start_time))

<|endoftext|>


<commit_before>import asyncio
import pytest
import unittest
from unittest import mock

from kafka.common import (KafkaError, ConnectionError,
                          NodeNotReadyError, UnrecognizedBrokerVersion)
from kafka.protocol.metadata import MetadataRequest, MetadataResponse

from aiokafka.client import AIOKafkaClient
from aiokafka.conn import AIOKafkaConnection
from ._testutil import KafkaIntegrationTestCase, run_until_complete

NO_ERROR = 0
UNKNOWN_TOPIC_OR_PARTITION = 3
NO_LEADER = 5
REPLICA_NOT_AVAILABLE = 9

@pytest.mark.usefixtures('setup_test_class')
class TestAIOKafkaClient(unittest.TestCase):

    def test_init_with_list(self):
        client = AIOKafkaClient(
            loop=self.loop,
            bootstrap_servers=['kafka01:9092', 'kafka02:9092', 'kafka03:9092'])
        self.assertEqual(
            '<AIOKafkaClient client_id=aiokafka-0.0.1>', client.__repr__())
        self.assertEqual(sorted({'kafka01': 9092,
                                 'kafka02': 9092,
                                 'kafka03': 9092}.items()),
                         sorted(client.hosts))

        node = client.get_random_node()
        self.assertEqual(node, None)  # unknown cluster metadata

    def test_init_with_csv(self):
        client = AIOKafkaClient(
            loop=self.loop,
            bootstrap_servers='kafka01:9092,kafka02:9092,kafka03:9092')

        self.assertEqual(sorted({'kafka01': 9092,
                                 'kafka02': 9092,
                                 'kafka03': 9092}.items()),
                         sorted(client.hosts))

    def test_load_metadata(self):
        brokers = [
            (0, 'broker_1', 4567),
            (1, 'broker_2', 5678)
        ]

        topics = [
            (NO_ERROR, 'topic_1', [
                (NO_ERROR, 0, 1, [1, 2], [1, 2])
            ]),
            (NO_ERROR, 'topic_2', [
                (NO_LEADER, 0, -1, [], []),
                (NO_LEADER, 1, 1, [], []),
            ]),
            (NO_LEADER, 'topic_no_partitions', []),
            (UNKNOWN_TOPIC_OR_PARTITION, 'topic_unknown', []),
            (NO_ERROR, 'topic_3', [
                (NO_ERROR, 0, 0, [0, 1], [0, 1]),
                (NO_ERROR, 1, 1, [1, 0], [1, 0]),
                (NO_ERROR, 2, 0, [0, 1], [0, 1])
            ]),
            (NO_ERROR, 'topic_4', [
                (NO_ERROR, 0, 0, [0, 1], [0, 1]),
                (REPLICA_NOT_AVAILABLE, 1, 1, [1, 0], [1, 0]),
            ])
        ]

        @asyncio.coroutine
        def send(request_id):
            return MetadataResponse(brokers, topics)

        mocked_conns = {0: mock.MagicMock()}
        mocked_conns[0].send.side_effect = send
        client = AIOKafkaClient(loop=self.loop,
                                bootstrap_servers=['broker_1:4567'])
        task = asyncio.async(client._md_synchronizer(), loop=self.loop)
        client._conns = mocked_conns
        client.cluster.update_metadata(MetadataResponse(brokers[:1], []))

        self.loop.run_until_complete(client.force_metadata_update())
        task.cancel()

        md = client.cluster
        c_brokers = md.brokers()
        self.assertEqual(len(c_brokers), 2)
        self.assertEqual(sorted(brokers), sorted(list(c_brokers)))
        c_topics = md.topics()
        self.assertEqual(len(c_topics), 4)
        self.assertEqual(md.partitions_for_topic('topic_1'), set([0]))
        self.assertEqual(md.partitions_for_topic('topic_2'), set([0, 1]))
        self.assertEqual(md.partitions_for_topic('topic_3'), set([0, 1, 2]))
        self.assertEqual(md.partitions_for_topic('topic_4'), set([0, 1]))
        self.assertEqual(
            md.available_partitions_for_topic('topic_2'), set([1]))

        mocked_conns[0].connected.return_value = False
        is_ready = self.loop.run_until_complete(client.ready(0))
        self.assertEqual(is_ready, False)
        is_ready = self.loop.run_until_complete(client.ready(1))
        self.assertEqual(is_ready, False)
        self.assertEqual(mocked_conns, {})

        with self.assertRaises(NodeNotReadyError):
            self.loop.run_until_complete(client.send(0, None))

class TestKafkaClientIntegration(KafkaIntegrationTestCase):

    @run_until_complete
    def test_bootstrap(self):
        client = AIOKafkaClient(loop=self.loop,
                                bootstrap_servers='0.42.42.42:444')
        with self.assertRaises(ConnectionError):
            yield from client.bootstrap()

        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        yield from client.bootstrap()
        yield from self.wait_topic(client, 'test_topic')

        metadata = yield from client.fetch_all_metadata()
        self.assertTrue('test_topic' in metadata.topics())

        client.set_topics(['t2', 't3'])
        client.set_topics(['t2', 't3'])  # should be ignored
        client.add_topic('t2')  # shold be ignored
        # bootstrap again -- no error expected
        yield from client.bootstrap()
        yield from client.close()

    @run_until_complete
    def test_failed_bootstrap(self):
        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        with mock.patch.object(AIOKafkaConnection, 'send') as mock_send:
            mock_send.side_effect = KafkaError('some kafka error')
            with self.assertRaises(ConnectionError):
                yield from client.bootstrap()

    @run_until_complete
    def test_send_request(self):
        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        yield from client.bootstrap()
        node_id = client.get_random_node()
        resp = yield from client.send(node_id, MetadataRequest([]))
        self.assertTrue(isinstance(resp, MetadataResponse))

    @run_until_complete
    def test_check_version(self):
        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        yield from client.bootstrap()
        ver = yield from client.check_version()
        self.assertTrue('0.' in ver)
        yield from self.wait_topic(client, 'some_test_topic')
        ver2 = yield from client.check_version()
        self.assertEqual(ver, ver2)
        ver2 = yield from client.check_version(client.get_random_node())
        self.assertEqual(ver, ver2)

        with mock.patch.object(
                AIOKafkaConnection, 'send') as mocked:
            mocked.side_effect = KafkaError('mocked exception')
            with self.assertRaises(UnrecognizedBrokerVersion):
                yield from client.check_version(client.get_random_node())

        client._get_conn = asyncio.coroutine(lambda _: None)
        with self.assertRaises(ConnectionError):
            yield from client.check_version()

    @run_until_complete
    def test_metadata_synchronizer(self):
        client = AIOKafkaClient(
            loop=self.loop,
            bootstrap_servers=self.hosts,
            metadata_max_age_ms=100)

        with mock.patch.object(
                AIOKafkaClient, '_metadata_update') as mocked:
            @asyncio.coroutine
            def dummy(*d, **kw):
                client.cluster.failed_update(None)
            mocked.side_effect = dummy

            yield from client.bootstrap()
            yield from asyncio.sleep(0.15, loop=self.loop)
            yield from client.close()

            self.assertNotEqual(
                len(client._metadata_update.mock_calls), 0)

    @run_until_complete
    def test_metadata_update_fail(self):
        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        yield from client.bootstrap()

        with mock.patch.object(
                AIOKafkaConnection, 'send') as mocked:
            mocked.side_effect = KafkaError('mocked exception')

            updated = yield from client.force_metadata_update()

            self.assertEqual(updated, False)

            with self.assertRaises(KafkaError):
                yield from client.fetch_all_metadata()
<commit_msg>Fix version related test<commit_after>import asyncio
import pytest
import unittest
from unittest import mock

from kafka.common import (KafkaError, ConnectionError,
                          NodeNotReadyError, UnrecognizedBrokerVersion)
from kafka.protocol.metadata import MetadataRequest, MetadataResponse

from aiokafka.client import AIOKafkaClient
from aiokafka.conn import AIOKafkaConnection
from ._testutil import KafkaIntegrationTestCase, run_until_complete

NO_ERROR = 0
UNKNOWN_TOPIC_OR_PARTITION = 3
NO_LEADER = 5
REPLICA_NOT_AVAILABLE = 9

@pytest.mark.usefixtures('setup_test_class')
class TestAIOKafkaClient(unittest.TestCase):

    def test_init_with_list(self):
        client = AIOKafkaClient(
            loop=self.loop,
            bootstrap_servers=['kafka01:9092', 'kafka02:9092', 'kafka03:9092'])
        self.assertEqual(
            '<AIOKafkaClient client_id=aiokafka-0.1.0>', client.__repr__())
        self.assertEqual(sorted({'kafka01': 9092,
                                 'kafka02': 9092,
                                 'kafka03': 9092}.items()),
                         sorted(client.hosts))

        node = client.get_random_node()
        self.assertEqual(node, None)  # unknown cluster metadata

    def test_init_with_csv(self):
        client = AIOKafkaClient(
            loop=self.loop,
            bootstrap_servers='kafka01:9092,kafka02:9092,kafka03:9092')

        self.assertEqual(sorted({'kafka01': 9092,
                                 'kafka02': 9092,
                                 'kafka03': 9092}.items()),
                         sorted(client.hosts))

    def test_load_metadata(self):
        brokers = [
            (0, 'broker_1', 4567),
            (1, 'broker_2', 5678)
        ]

        topics = [
            (NO_ERROR, 'topic_1', [
                (NO_ERROR, 0, 1, [1, 2], [1, 2])
            ]),
            (NO_ERROR, 'topic_2', [
                (NO_LEADER, 0, -1, [], []),
                (NO_LEADER, 1, 1, [], []),
            ]),
            (NO_LEADER, 'topic_no_partitions', []),
            (UNKNOWN_TOPIC_OR_PARTITION, 'topic_unknown', []),
            (NO_ERROR, 'topic_3', [
                (NO_ERROR, 0, 0, [0, 1], [0, 1]),
                (NO_ERROR, 1, 1, [1, 0], [1, 0]),
                (NO_ERROR, 2, 0, [0, 1], [0, 1])
            ]),
            (NO_ERROR, 'topic_4', [
                (NO_ERROR, 0, 0, [0, 1], [0, 1]),
                (REPLICA_NOT_AVAILABLE, 1, 1, [1, 0], [1, 0]),
            ])
        ]

        @asyncio.coroutine
        def send(request_id):
            return MetadataResponse(brokers, topics)

        mocked_conns = {0: mock.MagicMock()}
        mocked_conns[0].send.side_effect = send
        client = AIOKafkaClient(loop=self.loop,
                                bootstrap_servers=['broker_1:4567'])
        task = asyncio.async(client._md_synchronizer(), loop=self.loop)
        client._conns = mocked_conns
        client.cluster.update_metadata(MetadataResponse(brokers[:1], []))

        self.loop.run_until_complete(client.force_metadata_update())
        task.cancel()

        md = client.cluster
        c_brokers = md.brokers()
        self.assertEqual(len(c_brokers), 2)
        self.assertEqual(sorted(brokers), sorted(list(c_brokers)))
        c_topics = md.topics()
        self.assertEqual(len(c_topics), 4)
        self.assertEqual(md.partitions_for_topic('topic_1'), set([0]))
        self.assertEqual(md.partitions_for_topic('topic_2'), set([0, 1]))
        self.assertEqual(md.partitions_for_topic('topic_3'), set([0, 1, 2]))
        self.assertEqual(md.partitions_for_topic('topic_4'), set([0, 1]))
        self.assertEqual(
            md.available_partitions_for_topic('topic_2'), set([1]))

        mocked_conns[0].connected.return_value = False
        is_ready = self.loop.run_until_complete(client.ready(0))
        self.assertEqual(is_ready, False)
        is_ready = self.loop.run_until_complete(client.ready(1))
        self.assertEqual(is_ready, False)
        self.assertEqual(mocked_conns, {})

        with self.assertRaises(NodeNotReadyError):
            self.loop.run_until_complete(client.send(0, None))

class TestKafkaClientIntegration(KafkaIntegrationTestCase):

    @run_until_complete
    def test_bootstrap(self):
        client = AIOKafkaClient(loop=self.loop,
                                bootstrap_servers='0.42.42.42:444')
        with self.assertRaises(ConnectionError):
            yield from client.bootstrap()

        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        yield from client.bootstrap()
        yield from self.wait_topic(client, 'test_topic')

        metadata = yield from client.fetch_all_metadata()
        self.assertTrue('test_topic' in metadata.topics())

        client.set_topics(['t2', 't3'])
        client.set_topics(['t2', 't3'])  # should be ignored
        client.add_topic('t2')  # shold be ignored
        # bootstrap again -- no error expected
        yield from client.bootstrap()
        yield from client.close()

    @run_until_complete
    def test_failed_bootstrap(self):
        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        with mock.patch.object(AIOKafkaConnection, 'send') as mock_send:
            mock_send.side_effect = KafkaError('some kafka error')
            with self.assertRaises(ConnectionError):
                yield from client.bootstrap()

    @run_until_complete
    def test_send_request(self):
        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        yield from client.bootstrap()
        node_id = client.get_random_node()
        resp = yield from client.send(node_id, MetadataRequest([]))
        self.assertTrue(isinstance(resp, MetadataResponse))

    @run_until_complete
    def test_check_version(self):
        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        yield from client.bootstrap()
        ver = yield from client.check_version()
        self.assertTrue('0.' in ver)
        yield from self.wait_topic(client, 'some_test_topic')
        ver2 = yield from client.check_version()
        self.assertEqual(ver, ver2)
        ver2 = yield from client.check_version(client.get_random_node())
        self.assertEqual(ver, ver2)

        with mock.patch.object(
                AIOKafkaConnection, 'send') as mocked:
            mocked.side_effect = KafkaError('mocked exception')
            with self.assertRaises(UnrecognizedBrokerVersion):
                yield from client.check_version(client.get_random_node())

        client._get_conn = asyncio.coroutine(lambda _: None)
        with self.assertRaises(ConnectionError):
            yield from client.check_version()

    @run_until_complete
    def test_metadata_synchronizer(self):
        client = AIOKafkaClient(
            loop=self.loop,
            bootstrap_servers=self.hosts,
            metadata_max_age_ms=100)

        with mock.patch.object(
                AIOKafkaClient, '_metadata_update') as mocked:
            @asyncio.coroutine
            def dummy(*d, **kw):
                client.cluster.failed_update(None)
            mocked.side_effect = dummy

            yield from client.bootstrap()
            yield from asyncio.sleep(0.15, loop=self.loop)
            yield from client.close()

            self.assertNotEqual(
                len(client._metadata_update.mock_calls), 0)

    @run_until_complete
    def test_metadata_update_fail(self):
        client = AIOKafkaClient(loop=self.loop, bootstrap_servers=self.hosts)
        yield from client.bootstrap()

        with mock.patch.object(
                AIOKafkaConnection, 'send') as mocked:
            mocked.side_effect = KafkaError('mocked exception')

            updated = yield from client.force_metadata_update()

            self.assertEqual(updated, False)

            with self.assertRaises(KafkaError):
                yield from client.fetch_all_metadata()
<|endoftext|>


<commit_before>import nose
from nose.tools import raises
import utils
utils.setup_pythonpath()
import rez.config
from rez.config import Resolver
from rez.exceptions import PkgsUnresolvedError, PkgConfigNotResolvedError, PkgConflictError, PkgNotFoundError
from rez.public_enums import RESOLVE_MODE_LATEST, RESOLVE_MODE_EARLIEST
from rez.filesys import _g_os_pkg as OS_PKG
from rez.filesys import _g_arch_pkg as ARCH_PKG

def check_basic_resolve(pkgs, assertions,
                        resolver_args=dict(resolve_mode=RESOLVE_MODE_LATEST),
                        resolve_args={}):
    resolver = rez.config.Resolver(**resolver_args)
    result = resolver.resolve(pkgs, **resolve_args)
    # TODO: reset cached resolves
    assert_resolve_result(result, [OS_PKG] + assertions + [ARCH_PKG])

def assert_resolve_result(result, assertions):
    assert result is not None
    pkg_ress, commands, dot_graph, num_fails = result

    res = [p.short_name() for p in pkg_ress]
    assert res == assertions, res

class ResolveBaseTest(utils.BaseTest):
    def setUp(self):
        self.cleanup()
        self.add_packages()
        self.make_packages()

    def add_packages(self):
        # real world examples are so much easier to follow
        with self.add_package('python-2.7.4', local=True) as pkg:
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]

        with self.add_package('python-2.6.4') as pkg:
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]

        with self.add_package('python-2.6.1') as pkg:
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]

        with self.add_package('maya-2012') as pkg:
            pkg.requires = ['python-2.6']
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]

        with self.add_package('maya-2013') as pkg:
            pkg.requires = ['python-2.6']
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]

        with self.add_package('maya-2014') as pkg:
            pkg.requires = ['python-2.7']
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]

        with self.add_package('nuke-7.1.2') as pkg:
            pkg.requires = ['python-2.6']

        with self.add_package('arnold-4.0.16.0') as pkg:
            pkg.requires = ['python']
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]

        with self.add_package('mtoa-0.25.0') as pkg:
            #pkg.requires = ['arnold-4.0.16']
            pkg.variants = [['platform-linux', 'maya-2014', 'arnold-4.0.15+'],
                            ['platform-linux', 'maya-2013', 'arnold-4.0.15+'],
                            ['platform-darwin', 'maya-2014', 'arnold-4.0.15+'],
                            ['platform-darwin', 'maya-2013', 'arnold-4.0.15+']]

        self.add_package('platform-linux')
        self.add_package('platform-darwin')

        self.add_package('arch-x86_64')
        self.add_package('arch-i386')

        # versionless
        with self.add_package('site') as pkg:
            pkg.requires = ['maya', 'nuke-7']

class TestResolve(ResolveBaseTest):
    def test_latest(self):
        for ins, outs in [
                          (['python'],
                           ['python-2.7.4']),
                          (['python-2.6'],
                           ['python-2.6.4']),
                          (['maya'],
                           ['python-2.7.4', 'maya-2014']),
                          (['maya', 'python-2.6'],
                           ['python-2.6.4', 'maya-2013']),
                          (['maya', 'nuke-7'],
                           ['python-2.6.4', 'nuke-7.1.2', 'maya-2013']),
                          (['site'],
                           ['python-2.6.4', 'nuke-7.1.2', 'maya-2013', 'site']),
                          (['nuke-7'],
                           ['python-2.6.4', 'nuke-7.1.2']),
                          (['mtoa'],
                           ['python-2.7.4', 'maya-2014', 'arnold-4.0.16.0', 'mtoa-0.25.0']),
                          ([OS_PKG],
                           [])
                          ]:
            yield check_basic_resolve, ins, outs, dict(resolve_mode=RESOLVE_MODE_LATEST)
            yield check_basic_resolve, ins, outs, dict(resolve_mode=RESOLVE_MODE_LATEST,
                                                       assume_dt=True)

    def test_earliest(self):
        for ins, outs in [(['python'],
                           ['python-2.6.1']),
                          (['python-2.6'],
                           ['python-2.6.1']),
                          (['maya'],
                           ['python-2.6.1', 'maya-2012']),
                          (['maya', 'python-2.6'],
                           ['python-2.6.1', 'maya-2012']),
                          (['maya', 'nuke'],
                           ['python-2.6.1', 'nuke-7.1.2', 'maya-2012']),
                          ]:
            yield check_basic_resolve, ins, outs, dict(resolve_mode=RESOLVE_MODE_EARLIEST)

    def test_failures(self):
        for ins, exc in [(['python-2.7', 'python-2.6'], PkgConflictError), # straight conflict
                         (['nuke-6'], PkgNotFoundError), # does not exist
                         # I dont understand the practical difference between
                         # PkgsUnresolvedError and PkgConfigNotResolvedError
                         (['maya-2014', 'nuke-7'], PkgConfigNotResolvedError),
                         (['maya-2014', 'nuke-7+'], PkgConfigNotResolvedError),
                         ]:
            # `raises` is a decorator that returns a modified test function that
            # passes the test if the exception is raised
            yield raises(exc)(check_basic_resolve), ins, None

class TestCommands(TestResolve):
    PYTHON_COMMANDS = '''
REZ_PYTHON_MAJOR_VERSION = '{version.part(1)}'
REZ_PYTHON_MINOR_VERSION = '{version.part(2)}'

if machine.os == 'Linux':
    PYTHON_DIR = '/usr/local/python-{version}'
    PATH.prepend('$PYTHON_DIR/bin')
elif machine.os == 'Darwin':
    PYTHON_DIR = '/usr/local/python-{version}'
    PATH.prepend('$PYTHON_DIR/Python.framework/Versions/{version.thru(2)}/bin')
else:
    PYTHON_DIR = 'C:/Python{version.part(1)}{version.part(2)}'
    PATH.prepend('$PYTHON_DIR')
    PATH.prepend('$PYTHON_DIR/Scripts')'''

    def add_packages(self):
        super(TestCommands, self).add_packages()
        # overrides:
        with self.add_package('python-2.7.4', local=True) as pkg:
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]
            # new style:
            pkg.commands = self.PYTHON_COMMANDS

        with self.add_package('python-2.6.4') as pkg:
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]
            # new style:
            pkg.commands = self.PYTHON_COMMANDS

        with self.add_package('arnold-4.0.16.0') as pkg:
            pkg.requires = ['python']
            # old-style:
            pkg.commands = ['export CMAKE_MODULE_PATH=!ROOT!/cmake:$CMAKE_MODULE_PATH'
                            'export ARNOLD_HOME=/usr/local/solidAngle/arnold-!VERSION!'
                            'export PATH=$ARNOLD_HOME/bin:$PATH'
                            'export PYTHONPATH=$ARNOLD_HOME/python:$PYTHONPATH']

    def test_commands(self):
        for ins, outs in [
                          (['python'],
                           ['python-2.7.4']),
                          (['python-2.6'],
                           ['python-2.6.4']),
                          (['mtoa'],
                           ['python-2.7.4', 'maya-2014', 'arnold-4.0.16.0', 'mtoa-0.25.0']),
                          ]:
            yield check_basic_resolve, ins, outs

if __name__ == '__main__':
    nose.main()<commit_msg>Add a test package with 2 dimensional variants.  Add tools to packages for future testing of wrappers.<commit_after>import nose
from nose.tools import raises
import utils
utils.setup_pythonpath()
import rez.config
from rez.config import Resolver
from rez.exceptions import PkgsUnresolvedError, PkgConfigNotResolvedError, PkgConflictError, PkgNotFoundError
from rez.public_enums import RESOLVE_MODE_LATEST, RESOLVE_MODE_EARLIEST
from rez.filesys import _g_os_pkg as OS_PKG
from rez.filesys import _g_arch_pkg as ARCH_PKG

def check_basic_resolve(pkgs, assertions,
                        resolver_args=dict(resolve_mode=RESOLVE_MODE_LATEST),
                        resolve_args={}):
    resolver = rez.config.Resolver(**resolver_args)
    result = resolver.resolve(pkgs, **resolve_args)
    # TODO: reset cached resolves
    assert_resolve_result(result, [OS_PKG] + assertions + [ARCH_PKG])

def assert_resolve_result(result, assertions):
    assert result is not None
    pkg_ress, commands, dot_graph, num_fails = result

    res = [p.short_name() for p in pkg_ress]
    assert res == assertions, res

class ResolveBaseTest(utils.BaseTest):
    def setUp(self):
        self.cleanup()
        self.add_packages()
        self.make_packages()

    def add_packages(self):
        # real world examples are so much easier to follow
        with self.add_package('python-2.7.4', local=True) as pkg:
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]

        with self.add_package('python-2.6.4') as pkg:
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]

        with self.add_package('python-2.6.1') as pkg:
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]

        with self.add_package('mercurial-3.0') as pkg:
            pkg.variants = [['platform-linux', 'python-2.7'],
                            ['platform-linux', 'python-2.6'],
                            ['platform-darwin', 'python-2.7']]

        with self.add_package('maya-2012') as pkg:
            pkg.requires = ['python-2.6']
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]

        with self.add_package('maya-2013') as pkg:
            pkg.requires = ['python-2.6']
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]
            pkg.tools = ['maya', 'mayapy']

        with self.add_package('maya-2014') as pkg:
            pkg.requires = ['python-2.7']
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]
            pkg.tools = ['maya', 'mayapy']

        with self.add_package('nuke-7.1.2') as pkg:
            pkg.requires = ['python-2.6']
            pkg.tools = ['Nuke']

        with self.add_package('arnold-4.0.16.0') as pkg:
            pkg.requires = ['python']
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]
            pkg.tools = ['kick']

        with self.add_package('mtoa-0.25.0') as pkg:
            #pkg.requires = ['arnold-4.0.16']
            pkg.variants = [['platform-linux', 'maya-2014', 'arnold-4.0.15+'],
                            ['platform-linux', 'maya-2013', 'arnold-4.0.15+'],
                            ['platform-darwin', 'maya-2014', 'arnold-4.0.15+'],
                            ['platform-darwin', 'maya-2013', 'arnold-4.0.15+']]

        self.add_package('platform-linux')
        self.add_package('platform-darwin')

        self.add_package('arch-x86_64')
        self.add_package('arch-i386')

        # versionless
        with self.add_package('site') as pkg:
            pkg.requires = ['maya', 'nuke-7']

class TestResolve(ResolveBaseTest):
    def test_latest(self):
        for ins, outs in [
                          (['python'],
                           ['python-2.7.4']),
                          (['python-2.6'],
                           ['python-2.6.4']),
                          (['maya'],
                           ['python-2.7.4', 'maya-2014']),
                          (['maya', 'python-2.6'],
                           ['python-2.6.4', 'maya-2013']),
                          (['maya', 'nuke-7'],
                           ['python-2.6.4', 'nuke-7.1.2', 'maya-2013']),
                          (['site'],
                           ['python-2.6.4', 'nuke-7.1.2', 'maya-2013', 'site']),
                          (['nuke-7'],
                           ['python-2.6.4', 'nuke-7.1.2']),
                          (['mtoa'],
                           ['python-2.7.4', 'maya-2014', 'arnold-4.0.16.0', 'mtoa-0.25.0']),
                          (['python', 'mercurial'],
                           ['python-2.7.4', 'mercurial-3.0']),
                          ([OS_PKG],
                           [])
                          ]:
            yield check_basic_resolve, ins, outs, dict(resolve_mode=RESOLVE_MODE_LATEST)
            yield check_basic_resolve, ins, outs, dict(resolve_mode=RESOLVE_MODE_LATEST,
                                                       assume_dt=True)

    def test_earliest(self):
        for ins, outs in [(['python'],
                           ['python-2.6.1']),
                          (['python-2.6'],
                           ['python-2.6.1']),
                          (['maya'],
                           ['python-2.6.1', 'maya-2012']),
                          (['maya', 'python-2.6'],
                           ['python-2.6.1', 'maya-2012']),
                          (['maya', 'nuke'],
                           ['python-2.6.1', 'nuke-7.1.2', 'maya-2012']),
                          ]:
            yield check_basic_resolve, ins, outs, dict(resolve_mode=RESOLVE_MODE_EARLIEST)

    def test_failures(self):
        for ins, exc in [(['python-2.7', 'python-2.6'], PkgConflictError), # straight conflict
                         (['nuke-6'], PkgNotFoundError), # does not exist
                         # I dont understand the practical difference between
                         # PkgsUnresolvedError and PkgConfigNotResolvedError
                         (['maya-2014', 'nuke-7'], PkgConfigNotResolvedError),
                         (['maya-2014', 'nuke-7+'], PkgConfigNotResolvedError),
                         ]:
            # `raises` is a decorator that returns a modified test function that
            # passes the test if the exception is raised
            yield raises(exc)(check_basic_resolve), ins, None

class TestCommands(TestResolve):
    PYTHON_COMMANDS = '''
REZ_PYTHON_MAJOR_VERSION = '{version.part(1)}'
REZ_PYTHON_MINOR_VERSION = '{version.part(2)}'

if machine.os == 'Linux':
    PYTHON_DIR = '/usr/local/python-{version}'
    PATH.prepend('$PYTHON_DIR/bin')
elif machine.os == 'Darwin':
    PYTHON_DIR = '/usr/local/python-{version}'
    PATH.prepend('$PYTHON_DIR/Python.framework/Versions/{version.thru(2)}/bin')
else:
    PYTHON_DIR = 'C:/Python{version.part(1)}{version.part(2)}'
    PATH.prepend('$PYTHON_DIR')
    PATH.prepend('$PYTHON_DIR/Scripts')'''

    def add_packages(self):
        super(TestCommands, self).add_packages()
        # overrides:
        with self.add_package('python-2.7.4', local=True) as pkg:
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]
            # new style:
            pkg.commands = self.PYTHON_COMMANDS

        with self.add_package('python-2.6.4') as pkg:
            pkg.variants = [['platform-linux'],
                            ['platform-darwin']]
            # new style:
            pkg.commands = self.PYTHON_COMMANDS

        with self.add_package('arnold-4.0.16.0') as pkg:
            pkg.requires = ['python']
            # old-style:
            pkg.commands = ['export CMAKE_MODULE_PATH=!ROOT!/cmake:$CMAKE_MODULE_PATH'
                            'export ARNOLD_HOME=/usr/local/solidAngle/arnold-!VERSION!'
                            'export PATH=$ARNOLD_HOME/bin:$PATH'
                            'export PYTHONPATH=$ARNOLD_HOME/python:$PYTHONPATH']

    def test_commands(self):
        for ins, outs in [
                          (['python'],
                           ['python-2.7.4']),
                          (['python-2.6'],
                           ['python-2.6.4']),
                          (['mtoa'],
                           ['python-2.7.4', 'maya-2014', 'arnold-4.0.16.0', 'mtoa-0.25.0']),
                          ]:
            yield check_basic_resolve, ins, outs

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