mirror of https://github.com/celery/kombu.git
247 lines
7.1 KiB
Python
247 lines
7.1 KiB
Python
from __future__ import annotations
|
|
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
from kombu import Connection, Producer, pools
|
|
from kombu.connection import ConnectionPool
|
|
from kombu.utils.collections import eqhash
|
|
|
|
|
|
class test_ProducerPool:
|
|
Pool = pools.ProducerPool
|
|
|
|
class MyPool(pools.ProducerPool):
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
self.instance = Mock()
|
|
super().__init__(*args, **kwargs)
|
|
|
|
def Producer(self, connection):
|
|
return self.instance
|
|
|
|
def setup(self):
|
|
self.connections = Mock()
|
|
self.pool = self.Pool(self.connections, limit=10)
|
|
|
|
def test_close_resource(self):
|
|
self.pool.close_resource(Mock(name='resource'))
|
|
|
|
def test_releases_connection_when_Producer_raises(self):
|
|
self.pool.Producer = Mock()
|
|
self.pool.Producer.side_effect = IOError()
|
|
acq = self.pool._acquire_connection = Mock()
|
|
conn = acq.return_value = Mock()
|
|
with pytest.raises(IOError):
|
|
self.pool.create_producer()
|
|
conn.release.assert_called_with()
|
|
|
|
def test_prepare_release_connection_on_error(self):
|
|
pp = Mock()
|
|
p = pp.return_value = Mock()
|
|
p.revive.side_effect = IOError()
|
|
acq = self.pool._acquire_connection = Mock()
|
|
conn = acq.return_value = Mock()
|
|
p._channel = None
|
|
with pytest.raises(IOError):
|
|
self.pool.prepare(pp)
|
|
conn.release.assert_called_with()
|
|
|
|
def test_release_releases_connection(self):
|
|
p = Mock()
|
|
p.__connection__ = Mock()
|
|
self.pool.release(p)
|
|
p.__connection__.release.assert_called_with()
|
|
p.__connection__ = None
|
|
self.pool.release(p)
|
|
|
|
def test_init(self):
|
|
assert self.pool.connections is self.connections
|
|
|
|
def test_Producer(self):
|
|
assert isinstance(self.pool.Producer(Mock()), Producer)
|
|
|
|
def test_acquire_connection(self):
|
|
self.pool._acquire_connection()
|
|
self.connections.acquire.assert_called_with(block=True)
|
|
|
|
def test_new(self):
|
|
promise = self.pool.new()
|
|
producer = promise()
|
|
assert isinstance(producer, Producer)
|
|
self.connections.acquire.assert_called_with(block=True)
|
|
|
|
def test_setup_unlimited(self):
|
|
pool = self.Pool(self.connections, limit=None)
|
|
pool.setup()
|
|
assert not pool._resource.queue
|
|
|
|
def test_setup(self):
|
|
assert len(self.pool._resource.queue) == self.pool.limit
|
|
|
|
first = self.pool._resource.get_nowait()
|
|
producer = first()
|
|
assert isinstance(producer, Producer)
|
|
|
|
def test_prepare(self):
|
|
connection = self.connections.acquire.return_value = Mock()
|
|
pool = self.MyPool(self.connections, limit=10)
|
|
pool.instance._channel = None
|
|
first = pool._resource.get_nowait()
|
|
producer = pool.prepare(first)
|
|
self.connections.acquire.assert_called()
|
|
producer.revive.assert_called_with(connection)
|
|
|
|
def test_prepare_channel_already_created(self):
|
|
self.connections.acquire.return_value = Mock()
|
|
pool = self.MyPool(self.connections, limit=10)
|
|
pool.instance._channel = Mock()
|
|
first = pool._resource.get_nowait()
|
|
self.connections.acquire.reset()
|
|
producer = pool.prepare(first)
|
|
producer.revive.assert_not_called()
|
|
|
|
def test_prepare_not_callable(self):
|
|
x = Producer(Mock)
|
|
self.pool.prepare(x)
|
|
|
|
def test_release(self):
|
|
p = Mock()
|
|
p.channel = Mock()
|
|
p.__connection__ = Mock()
|
|
self.pool.release(p)
|
|
p.__connection__.release.assert_called_with()
|
|
assert p.channel is None
|
|
|
|
|
|
class test_PoolGroup:
|
|
Group = pools.PoolGroup
|
|
|
|
class MyGroup(pools.PoolGroup):
|
|
|
|
def create(self, resource, limit):
|
|
return resource, limit
|
|
|
|
def test_interface_create(self):
|
|
g = self.Group()
|
|
with pytest.raises(NotImplementedError):
|
|
g.create(Mock(), 10)
|
|
|
|
def test_getitem_using_global_limit(self):
|
|
g = self.MyGroup(limit=pools.use_global_limit)
|
|
res = g['foo']
|
|
assert res == ('foo', pools.get_limit())
|
|
|
|
def test_getitem_using_custom_limit(self):
|
|
g = self.MyGroup(limit=102456)
|
|
res = g['foo']
|
|
assert res == ('foo', 102456)
|
|
|
|
def test_delitem(self):
|
|
g = self.MyGroup()
|
|
g['foo']
|
|
del(g['foo'])
|
|
assert 'foo' not in g
|
|
|
|
def test_Connections(self):
|
|
conn = Connection('memory://')
|
|
p = pools.connections[conn]
|
|
assert p
|
|
assert isinstance(p, ConnectionPool)
|
|
assert p.connection is conn
|
|
assert p.limit == pools.get_limit()
|
|
|
|
def test_Producers(self):
|
|
conn = Connection('memory://')
|
|
p = pools.producers[conn]
|
|
assert p
|
|
assert isinstance(p, pools.ProducerPool)
|
|
assert p.connections is pools.connections[conn]
|
|
assert p.limit == p.connections.limit
|
|
assert p.limit == pools.get_limit()
|
|
|
|
def test_all_groups(self):
|
|
conn = Connection('memory://')
|
|
pools.connections[conn]
|
|
|
|
assert list(pools._all_pools())
|
|
|
|
def test_reset(self):
|
|
pools.reset()
|
|
|
|
class MyGroup(dict):
|
|
clear_called = False
|
|
|
|
def clear(self):
|
|
self.clear_called = True
|
|
|
|
p1 = pools.connections['foo'] = Mock()
|
|
g1 = MyGroup()
|
|
pools._groups.append(g1)
|
|
|
|
pools.reset()
|
|
p1.force_close_all.assert_called_with()
|
|
assert g1.clear_called
|
|
|
|
p1 = pools.connections['foo'] = Mock()
|
|
p1.force_close_all.side_effect = KeyError()
|
|
pools.reset()
|
|
|
|
def test_set_limit(self):
|
|
pools.reset()
|
|
pools.set_limit(34576)
|
|
limit = pools.get_limit()
|
|
assert limit == 34576
|
|
|
|
conn = Connection('memory://')
|
|
pool = pools.connections[conn]
|
|
with pool.acquire():
|
|
pools.set_limit(limit + 1)
|
|
assert pools.get_limit() == limit + 1
|
|
limit = pools.get_limit()
|
|
with pytest.raises(RuntimeError):
|
|
pools.set_limit(limit - 1)
|
|
pools.set_limit(limit - 1, force=True)
|
|
assert pools.get_limit() == limit - 1
|
|
|
|
pools.set_limit(pools.get_limit())
|
|
|
|
def test_remove_limit(self):
|
|
conn = Connection('memory://')
|
|
pool = pools.connections[conn]
|
|
pool.limit = 10
|
|
with pool.acquire():
|
|
pool.limit = 0
|
|
|
|
|
|
class test_fun_PoolGroup:
|
|
|
|
def test_connections_behavior(self):
|
|
c1u = 'memory://localhost:123'
|
|
c2u = 'memory://localhost:124'
|
|
c1 = Connection(c1u)
|
|
c2 = Connection(c2u)
|
|
c3 = Connection(c1u)
|
|
|
|
assert eqhash(c1) != eqhash(c2)
|
|
assert eqhash(c1) == eqhash(c3)
|
|
|
|
c4 = Connection(c1u, transport_options={'confirm_publish': True})
|
|
assert eqhash(c3) != eqhash(c4)
|
|
|
|
p1 = pools.connections[c1]
|
|
p2 = pools.connections[c2]
|
|
p3 = pools.connections[c3]
|
|
|
|
assert p1 is not p2
|
|
assert p1 is p3
|
|
|
|
r1 = p1.acquire()
|
|
assert p1._dirty
|
|
assert p3._dirty
|
|
assert not p2._dirty
|
|
r1.release()
|
|
assert not p1._dirty
|
|
assert not p3._dirty
|