def setUp(self): prometheus.reset() assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 0 import ssl ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH) ssl_context.load_cert_chain('hpfeeds/tests/testcert.crt', 'hpfeeds/tests/testcert.key') self.sock = sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.bind(('127.0.0.1', 0)) self.port = sock.getsockname()[1] authenticator = Authenticator({ 'test': { 'secret': 'secret', 'subchans': ['test-chan'], 'pubchans': ['test-chan'], 'owner': 'some-owner', } }) self.server = Server(authenticator, sock=self.sock, ssl=ssl_context)
def setUp(self): setup_asyncio_reactor(self) prometheus.reset() self.loop = asyncio.get_event_loop() assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_client_connections') == 0 assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_connection_made') == 0 assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_connection_ready', {'ident': 'test'}) is None self.sock = sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.bind(('127.0.0.1', 0)) self.port = sock.getsockname()[1] authenticator = Authenticator({ 'test': { 'secret': 'secret', 'subchans': ['test-chan'], 'pubchans': ['test-chan'], 'owner': 'some-owner', } }) self.server = Server(authenticator, sock=self.sock)
class TestBrokerPrometheusEndpoint(unittest.TestCase): def setUp(self): prometheus.reset() authenticator = Authenticator({ 'test': { 'secret': 'secret', 'subchans': ['test-chan'], 'pubchans': ['test-chan'], 'owner': 'some-owner', } }) self.server = Server(authenticator, exporter='127.0.0.1:20001') self.server.add_endpoint_legacy('127.0.0.1:20000') def test_metrics_server(self): async def inner(): server_future = asyncio.ensure_future(self.server.serve_forever()) await self.server.when_started async with aiohttp.ClientSession() as session: async with session.get('http://127.0.0.1:20001/metrics') as resp: metrics = await resp.text() print(metrics) assert 'hpfeeds_broker_client_connections 0.0' in metrics assert 'hpfeeds_broker_connection_send_buffer_size{' not in metrics sock = socket.socket() sock.connect(('127.0.0.1', 20000)) async with aiohttp.ClientSession() as session: async with session.get('http://127.0.0.1:20001/metrics') as resp: metrics = await resp.text() print(metrics) assert 'hpfeeds_broker_client_connections 1.0' in metrics assert 'hpfeeds_broker_connection_send_buffer_size{' not in metrics sock.close() async with ClientSession('127.0.0.1', 20000, 'test', 'secret'): async with aiohttp.ClientSession() as session: async with session.get('http://127.0.0.1:20001/metrics') as resp: metrics = await resp.text() print(metrics) assert 'hpfeeds_broker_client_connections 1.0' in metrics assert 'hpfeeds_broker_connection_send_buffer_size{ident="test"} 0.0' in metrics server_future.cancel() await server_future asyncio.get_event_loop().run_until_complete(inner())
def setUp(self): authenticator = Authenticator({ 'test': { 'secret': 'secret', 'subchans': ['test-chan'], 'pubchans': ['test-chan'], 'owner': 'some-owner', } }) self.server = Server(authenticator) self.server.add_endpoint_str("tls:interface=127.0.0.1:port=0:cert=tests/testcert.crt:key=tests/testcert.key")
def setUp(self): prometheus.reset() assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_client_connections') == 0 self.sock = sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.bind(('127.0.0.1', 0)) self.port = sock.getsockname()[1] authenticator = self.make_authenticator() self.server = Server(authenticator) self.server.add_endpoint_test(sock)
def setUp(self): prometheus.reset() authenticator = Authenticator({ 'test': { 'secret': 'secret', 'subchans': ['test-chan'], 'pubchans': ['test-chan'], 'owner': 'some-owner', } }) self.server = Server(authenticator, exporter='127.0.0.1:20001') self.server.add_endpoint_legacy('127.0.0.1:20000')
def main(): parser = argparse.ArgumentParser(description='Run a hpfeeds broker') parser.add_argument('--bind', default='0.0.0.0:20000', action='store') parser.add_argument('--exporter', default='', action='store') parser.add_argument('--name', default='hpfeeds', action='store') parser.add_argument('--debug', default=False, action='store_true') parser.add_argument('--auth', default='mongodb', action='store') parser.add_argument('--mongo_host', default="mongodb", action='store') parser.add_argument('--mongo_port', default=27017, action='store') parser.add_argument('--tlscert', default=None, action='store') parser.add_argument('--tlskey', default=None, action='store') args = parser.parse_args() if (args.tlscert and not args.tlskey) or (args.tlskey and not args.tlscert): parser.error('Must specify --tlskey AND --tlscert') return logging.basicConfig( level=logging.DEBUG if args.debug else logging.INFO, ) if args.auth == 'mongodb': logging.info("Opening connection to MongoDB at %s:%s", args.mongo_host, args.mongo_port) auth = mongodb.Authenticator(args.mongo_host, args.mongo_port) elif args.auth == 'env': auth = env.Authenticator() else: auth = sqlite.Authenticator('sqlite.db') ssl_context = None if args.tlscert: ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH) ssl_context.load_cert_chain(args.tlscert, args.tlskey) broker = Server( auth=auth, bind=args.bind, exporter=args.exporter, name=args.name, ssl=ssl_context, ) return aiorun.run(broker.serve_forever())
def setUp(self): prometheus.reset() assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 0 self.sock = sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.bind(('127.0.0.1', 0)) self.port = sock.getsockname()[1] authenticator = Authenticator({ 'test': { 'secret': 'secret', 'subchans': ['test-chan'], 'pubchans': ['test-chan'], 'owner': 'some-owner', } }) self.server = Server(authenticator, sock=self.sock)
async def inner(): authenticator = Authenticator({ 'test': { 'secret': 'secret', 'subchans': ['test-chan'], 'pubchans': ['test-chan'], 'owner': 'some-owner', } }) self.server = Server(authenticator, sock=self.sock) self.log.debug('Starting server') future = asyncio.ensure_future(self.server.serve_forever()) self.log.debug('Awaiting test teardown') await self.server_future self.log.debug('Stopping test server') future.cancel() await future
def setUp(self): prometheus.reset() authenticator = Authenticator({ 'test': { 'secret': 'secret', 'subchans': ['test-chan'], 'pubchans': ['test-chan'], 'owner': 'some-owner', } }) self.server = Server(authenticator, bind='127.0.0.1:20000')
def main(): parser = argparse.ArgumentParser(description='Run a hpfeeds broker') parser.add_argument('--bind', default=None, action='store') parser.add_argument('--exporter', default='', action='store') parser.add_argument('--name', default='hpfeeds', action='store') parser.add_argument('--debug', default=False, action='store_true') parser.add_argument('--auth', default=None, action='append') parser.add_argument('--tlscert', default=None, action='store') parser.add_argument('--tlskey', default=None, action='store') parser.add_argument('-e', '--endpoint', default=None, action='append') args = parser.parse_args() if (args.tlscert and not args.tlskey) or (args.tlskey and not args.tlscert): parser.error('Must specify --tlskey AND --tlscert') return logging.basicConfig(level=logging.DEBUG if args.debug else logging.INFO, ) auth = multi.Authenticator() auths = args.auth if args.auth else ['sqlite'] for a in auths: try: auth.add(get_authenticator(a)) except ServerException as e: print(str(e)) sys.exit(1) broker = Server( auth=auth, exporter=args.exporter, name=args.name, ) if args.bind or not args.endpoint: bind = args.bind or '0.0.0.0:20000' broker.add_endpoint_legacy(bind, tlscert=args.tlscert, tlskey=args.tlskey) if args.endpoint: for endpoint in args.endpoint: broker.add_endpoint_str(endpoint) return aiorun.run(broker.serve_forever())
def main(): logging.basicConfig(level=logging.DEBUG if config.DEBUG else logging.INFO) s = Server() s.serve_forever() return 0
class TestClientIntegrationWithAioBroker(unittest.TestCase): def setUp(self): setup_asyncio_reactor(self) prometheus.reset() self.loop = asyncio.get_event_loop() assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_client_connections') == 0 assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_connection_made') == 0 assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_connection_ready', {'ident': 'test'}) is None self.sock = sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.bind(('127.0.0.1', 0)) self.port = sock.getsockname()[1] authenticator = Authenticator({ 'test': { 'secret': 'secret', 'subchans': ['test-chan'], 'pubchans': ['test-chan'], 'owner': 'some-owner', } }) self.server = Server(authenticator, sock=self.sock) def test_subscribe_and_publish(self): async def inner(): print('Starting server') server_future = asyncio.ensure_future(self.server.serve_forever()) print('Creating client service') client = ClientSessionService(f'tcp:127.0.0.1:{self.port}', 'test', 'secret') client.subscribe('test-chan') client.startService() # Wait till client connected await client.whenConnected.asFuture(self.loop) assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_client_connections') == 1 assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_connection_made') == 1 print('Publishing test message') client.publish('test-chan', b'test message') print('Waiting for read()') assert ('test', 'test-chan', b'test message') == await client.read().asFuture(self.loop) # We would test this after call to subscribe, but need to wait until sure server has processed command assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_subscriptions', {'ident': 'test', 'chan': 'test-chan'}) == 1 # This will only have incremented when server has processed auth message # Test can only reliably assert this is the case after reading a message assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_connection_ready', {'ident': 'test'}) == 1 print('Stopping client') await client.stopService().asFuture(self.loop) print('Stopping server') server_future.cancel() await server_future self.loop.run_until_complete(inner()) assert len(self.server.connections) == 0, 'Connection left dangling' assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_client_connections') == 0 assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_connection_lost', {'ident': 'test'}) == 1 # Closing should auto unsubscribe assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_subscriptions', {'ident': 'test', 'chan': 'test-chan'}) == 0 def test_late_subscribe_and_publish(self): async def inner(): print('Starting server') server_future = asyncio.ensure_future(self.server.serve_forever()) print('Creating client service') client = ClientSessionService(f'tcp:127.0.0.1:{self.port}', 'test', 'secret') client.startService() # Wait till client connected await client.whenConnected.asFuture(self.loop) assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_client_connections') == 1 assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_connection_made') == 1 # Subscribe to a new thing after connection is up client.subscribe('test-chan') print('Publishing test message') client.publish('test-chan', b'test message') print('Waiting for read()') assert ('test', 'test-chan', b'test message') == await client.read().asFuture(self.loop) # We would test this after call to subscribe, but need to wait until sure server has processed command assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_subscriptions', {'ident': 'test', 'chan': 'test-chan'}) == 1 # Unsubscribe while the connection is up client.unsubscribe('test-chan') # This will only have incremented when server has processed auth message # Test can only reliably assert this is the case after reading a message assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_connection_ready', {'ident': 'test'}) == 1 print('Stopping client') await client.stopService().asFuture(self.loop) print('Stopping server') server_future.cancel() await server_future self.loop.run_until_complete(inner()) assert len(self.server.connections) == 0, 'Connection left dangling' assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_client_connections') == 0 assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_connection_lost', {'ident': 'test'}) == 1 # Again, we should test this directly after calling unsubscribe(), but no ability to wait assert prometheus.REGISTRY.get_sample_value('hpfeeds_broker_subscriptions', {'ident': 'test', 'chan': 'test-chan'}) == 0
class TestClientIntegration(unittest.TestCase): log = logging.getLogger('hpfeeds.testserver') def _server_thread(self): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) self.server_future = loop.create_future() async def inner(): authenticator = Authenticator({ 'test': { 'secret': 'secret', 'subchans': ['test-chan'], 'pubchans': ['test-chan'], 'owner': 'some-owner', } }) self.server = Server(authenticator, sock=self.sock) self.log.debug('Starting server') future = asyncio.ensure_future(self.server.serve_forever()) self.log.debug('Awaiting test teardown') await self.server_future self.log.debug('Stopping test server') future.cancel() await future loop.run_until_complete(inner()) def setUp(self): prometheus.reset() assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 0 self.sock = sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.bind(('127.0.0.1', 0)) self.port = sock.getsockname()[1] self.server_thread = threading.Thread(target=self._server_thread, ) self.server_thread.start() def test_subscribe_and_publish(self): c = client.new('127.0.0.1', self.port, 'test', 'secret') c.subscribe('test-chan') c._subscribe() # If we have subscribed to a channel we should be able to see a # connection in monitoring assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 1 c.publish('test-chan', b'data') opcode, data = c._read_message() assert opcode == 3 assert readpublish(data) == ('test', 'test-chan', b'data') # We managed to publish a message - check this is reflected in stats assert 1 == prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_receive_publish_count', { 'ident': 'test', 'chan': 'test-chan' }) # If we managed to read a message from the broker then we must be subscribed # Check this is reflected in stats assert 1 == prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_subscriptions', { 'ident': 'test', 'chan': 'test-chan' }) self.log.debug('Stopping client') c.stop() self.log.debug('Closing client') c.close() def tearDown(self): self.log.debug('Cancelling future') self.server_future.set_result(None) self.log.debug('Waiting') self.server_thread.join() assert len(self.server.connections) == 0, 'Connection left dangling' assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 0 assert 0 == prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_subscriptions', { 'ident': 'test', 'chan': 'test-chan' })
class TestAsyncioClientIntegration(unittest.TestCase): log = logging.getLogger('hpfeeds.test_asyncio_client') def setUp(self): prometheus.reset() assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 0 self.sock = sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.bind(('127.0.0.1', 0)) self.port = sock.getsockname()[1] authenticator = Authenticator({ 'test': { 'secret': 'secret', 'subchans': ['test-chan'], 'pubchans': ['test-chan'], 'owner': 'some-owner', } }) self.server = Server(authenticator, sock=self.sock) def test_subscribe_and_publish(self): async def inner(): self.log.debug('Starting server') server_future = asyncio.ensure_future(self.server.serve_forever()) self.log.debug('Creating client service') client = ClientSession('127.0.0.1', self.port, 'test', 'secret') client.subscribe('test-chan') # Wait till client connected await client.when_connected assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 1 assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_made') == 1 self.log.debug('Publishing test message') client.publish('test-chan', b'test message') self.log.debug('Waiting for read()') assert ('test', 'test-chan', b'test message') == await client.read() # We would test this after call to subscribe, but need to wait until sure server has processed command assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_subscriptions', { 'ident': 'test', 'chan': 'test-chan' }) == 1 # This will only have incremented when server has processed auth message # Test can only reliably assert this is the case after reading a message assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_ready', {'ident': 'test'}) == 1 self.log.debug('Stopping client') await client.close() assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_send_buffer_fill', {'ident': 'test'}) == 12 assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_send_buffer_drain', {'ident': 'test'}) == 32 self.log.debug('Stopping server') server_future.cancel() await server_future asyncio.get_event_loop().run_until_complete(inner()) assert len(self.server.connections) == 0, 'Connection left dangling' assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 0 assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_lost', {'ident': 'test'}) == 1 # Closing should auto unsubscribe assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_subscriptions', { 'ident': 'test', 'chan': 'test-chan' }) == 0 def test_late_subscribe_and_publish(self): async def inner(): self.log.debug('Starting server') server_future = asyncio.ensure_future(self.server.serve_forever()) self.log.debug('Creating client service') client = ClientSession('127.0.0.1', self.port, 'test', 'secret') # Wait till client connected await client.when_connected assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 1 assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_made') == 1 # Subscribe to a new thing after connection is up client.subscribe('test-chan') self.log.debug('Publishing test message') client.publish('test-chan', b'test message') self.log.debug('Waiting for read()') assert ('test', 'test-chan', b'test message') == await client.read() # We would test this after call to subscribe, but need to wait until sure server has processed command assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_subscriptions', { 'ident': 'test', 'chan': 'test-chan' }) == 1 # Unsubscribe while the connection is up client.unsubscribe('test-chan') # FIXME: How to test that did anything! # This will only have incremented when server has processed auth message # Test can only reliably assert this is the case after reading a message assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_ready', {'ident': 'test'}) == 1 self.log.debug('Stopping client') await client.close() self.log.debug('Stopping server') server_future.cancel() await server_future asyncio.get_event_loop().run_until_complete(inner()) assert len(self.server.connections) == 0, 'Connection left dangling' assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 0 assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_lost', {'ident': 'test'}) == 1 # Again, we should test this directly after calling unsubscribe(), but no ability to wait assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_subscriptions', { 'ident': 'test', 'chan': 'test-chan' }) == 0 def test_late_subscribe_and_publish_async_with(self): async def inner(): self.log.debug('Starting server') server_future = asyncio.ensure_future(self.server.serve_forever()) self.log.debug('Creating client service') async with ClientSession('127.0.0.1', self.port, 'test', 'secret') as client: assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 1 assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_made') == 1 # Subscribe to a new thing after connection is up client.subscribe('test-chan') self.log.debug('Publishing test message') client.publish('test-chan', b'test message') self.log.debug('Waiting for read()') assert ('test', 'test-chan', b'test message') == await client.read() # We would test this after call to subscribe, but need to wait until sure server has processed command assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_subscriptions', { 'ident': 'test', 'chan': 'test-chan' }) == 1 # Unsubscribe while the connection is up client.unsubscribe('test-chan') # FIXME: How to test that did anything! # This will only have incremented when server has processed auth message # Test can only reliably assert this is the case after reading a message assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_ready', {'ident': 'test'}) == 1 self.log.debug('Stopping server') server_future.cancel() await server_future asyncio.get_event_loop().run_until_complete(inner()) assert len(self.server.connections) == 0, 'Connection left dangling' assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 0 assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_lost', {'ident': 'test'}) == 1 # Again, we should test this directly after calling unsubscribe(), but no ability to wait assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_subscriptions', { 'ident': 'test', 'chan': 'test-chan' }) == 0 def test_late_subscribe_and_publish_async_for(self): async def inner(): server_future = asyncio.ensure_future(self.server.serve_forever()) async with ClientSession('127.0.0.1', self.port, 'test', 'secret') as client: client.subscribe('test-chan') client.publish('test-chan', b'test message') async for ident, chan, payload in client: assert ident == 'test' assert chan == 'test-chan' assert payload == b'test message' break server_future.cancel() await server_future asyncio.get_event_loop().run_until_complete(inner()) assert len(self.server.connections) == 0, 'Connection left dangling' def test_late_subscribe_and_publish_for_async_iter(self): async def inner(): server_future = asyncio.ensure_future(self.server.serve_forever()) async def example_iter(): yield b'test message' async with ClientSession('127.0.0.1', self.port, 'test', 'secret') as client: client.subscribe('test-chan') await client.publish_async_iterable('test-chan', example_iter()) assert ('test', 'test-chan', b'test message') == await client.read() server_future.cancel() await server_future asyncio.get_event_loop().run_until_complete(inner()) assert len(self.server.connections) == 0, 'Connection left dangling'
class TestBrokerIntegration(unittest.TestCase): log = logging.getLogger('hpfeeds.test_broker_connection') def setUp(self): authenticator = Authenticator({ 'test': { 'secret': 'secret', 'subchans': ['test-chan'], 'pubchans': ['test-chan'], 'owner': 'some-owner', } }) self.server = Server(authenticator) self.server.add_endpoint_str( "tls:interface=127.0.0.1:port=0:cert=hpfeeds/tests/testcert.crt:key=hpfeeds/tests/testcert.key" ) def test_subscribe_and_publish(self): prometheus.reset() async def inner(): self.log.debug('Starting server') server_future = asyncio.ensure_future(self.server.serve_forever()) await self.server.when_started self.port = self.server.endpoints[0]['port'] import ssl ssl_context = ssl.create_default_context( ssl.Purpose.SERVER_AUTH, cafile='hpfeeds/tests/testcert.crt') ssl_context.check_hostname = False self.log.debug('Creating client service') client = ClientSession('127.0.0.1', self.port, 'test', 'secret', ssl=ssl_context) client.subscribe('test-chan') # Wait till client connected await client.when_connected assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 1 assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_made') == 1 self.log.debug('Publishing test message') client.publish('test-chan', b'test message') self.log.debug('Waiting for read()') assert ('test', 'test-chan', b'test message') == await client.read() # We would test this after call to subscribe, but need to wait until sure server has processed command assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_subscriptions', { 'ident': 'test', 'chan': 'test-chan' }) == 1 # This will only have incremented when server has processed auth message # Test can only reliably assert this is the case after reading a message assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_ready', {'ident': 'test'}) == 1 self.log.debug('Stopping client') await client.close() assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_send_buffer_fill', {'ident': 'test'}) == 12 assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_send_buffer_drain', {'ident': 'test'}) == 32 self.log.debug('Stopping server') server_future.cancel() await server_future asyncio.get_event_loop().run_until_complete(inner()) assert len(self.server.connections) == 0, 'Connection left dangling' assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_client_connections') == 0 assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_connection_lost', {'ident': 'test'}) == 1 # Closing should auto unsubscribe assert prometheus.REGISTRY.get_sample_value( 'hpfeeds_broker_subscriptions', { 'ident': 'test', 'chan': 'test-chan' }) == 0
class TestBrokerConnection(unittest.TestCase): def setUp(self): prometheus.reset() authenticator = Authenticator({ 'test': { 'secret': 'secret', 'subchans': ['test-chan'], 'pubchans': ['test-chan'], 'owner': 'some-owner', } }) self.server = Server(authenticator) self.server.add_endpoint_legacy('127.0.0.1:20000') def make_connection(self): transport = mock.Mock() transport.get_extra_info.side_effect = lambda name: ( '127.0.0.1', 80) if name == 'peername' else None connection = Connection(self.server) connection.connection_made(transport) return connection def test_sends_challenge(self): c = self.make_connection() assert parse(c.transport.write)[0][1][:-4] == b'\x07hpfeeds' def test_must_auth(self): c = self.make_connection() c.data_received(msgpublish('a', 'b', b'c')) assert parse(c.transport.write)[0][1][:-4] == b'\x07hpfeeds' assert parse(c.transport.write)[1][1] == b'First message was not AUTH' def test_auth_failure_wrong_secret(self): c = self.make_connection() name, rand = readinfo(parse(c.transport.write)[0][1]) c.data_received(msgauth(rand, 'test', 'secret2')) assert parse( c.transport.write)[1][1] == b'Authentication failed for test' def test_auth_failure_no_such_ident(self): c = self.make_connection() name, rand = readinfo(parse(c.transport.write)[0][1]) c.data_received(msgauth(rand, 'test2', 'secret')) assert parse( c.transport.write)[1][1] == b'Authentication failed for test2' def test_permission_to_sub(self): c = self.make_connection() name, rand = readinfo(parse(c.transport.write)[0][1]) c.data_received(msgauth(rand, 'test', 'secret')) c.data_received(msgsubscribe('test', 'test-chan2')) assert parse(c.transport.write)[1][ 1] == b'Authkey not allowed to sub here. ident=test, chan=test-chan2' def test_permission_to_pub(self): c = self.make_connection() name, rand = readinfo(parse(c.transport.write)[0][1]) c.data_received(msgauth(rand, 'test', 'secret')) c.data_received(msgpublish('test', 'test-chan2', b'c')) assert parse(c.transport.write)[1][ 1] == b'Authkey not allowed to pub here. ident=test, chan=test-chan2' def test_pub_ident_checked(self): c = self.make_connection() name, rand = readinfo(parse(c.transport.write)[0][1]) c.data_received(msgauth(rand, 'test', 'secret')) c.data_received(msgpublish('wrong-ident', 'test-chan2', b'c')) assert parse( c.transport.write )[1][1] == b'Invalid authkey in message, ident=wrong-ident' def test_auth_success(self): c = self.make_connection() name, rand = readinfo(parse(c.transport.write)[0][1]) c.data_received(msgauth(rand, 'test', 'secret')) c.data_received(msgsubscribe('test', 'test-chan')) c.data_received(msgpublish('test', 'test-chan', b'c')) assert readpublish(parse(c.transport.write)[1][1]) == ('test', 'test-chan', b'c') def test_multiple_subscribers(self): subscribers = [] for i in range(5): c = self.make_connection() name, rand = readinfo(parse(c.transport.write)[0][1]) c.data_received(msgauth(rand, 'test', 'secret')) c.data_received(msgsubscribe('test', 'test-chan')) subscribers.append(c) c = self.make_connection() name, rand = readinfo(parse(c.transport.write)[0][1]) c.data_received(msgauth(rand, 'test', 'secret')) c.data_received(msgpublish('test', 'test-chan', b'c')) for c in subscribers: msgs = parse(c.transport.write) assert readpublish(msgs[1][1]) == ('test', 'test-chan', b'c') def test_auth_unsubscribe(self): c = self.make_connection() name, rand = readinfo(parse(c.transport.write)[0][1]) c.data_received(msgauth(rand, 'test', 'secret')) c.data_received(msgsubscribe('test', 'test-chan')) c.data_received(msgpublish('test', 'test-chan', b'c')) c.data_received(msgunsubscribe('test', 'test-chan')) c.data_received(msgpublish('test', 'test-chan', b'c')) c.data_received(msgsubscribe('test', 'test-chan')) c.data_received(msgpublish('test', 'test-chan', b'c')) messages = parse(c.transport.write) for msg in messages[1:]: assert readpublish(msg[1]) == ('test', 'test-chan', b'c') # 1 auth and 2 publish assert len(messages) == 3