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)
Exemple #3
0
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)
Exemple #6
0
    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')
Exemple #7
0
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)
Exemple #9
0
        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
Exemple #10
0
    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')
Exemple #11
0
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())
Exemple #12
0
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
Exemple #14
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