def test_pattern(self): """Checks instance aware pattern construction.""" app_discovery = discovery.Discovery(None, 'appproid.foo', 'http') self.assertEqual('foo#*', app_discovery.pattern) app_discovery = discovery.Discovery(None, 'appproid.foo#1', 'http') self.assertEqual('foo#1', app_discovery.pattern)
def test_sync(self): """Checks event processing and state sync.""" zkclient = kazoo.client.KazooClient() app_discovery = discovery.Discovery(zkclient, 'appproid.foo.*', 'http') kazoo.client.KazooClient.get_children.return_value = [ 'foo.1#0:tcp:http', 'foo.2#0:tcp:http', 'foo.2#0:tcp:tcp', 'bar.1#0:tcp:http' ] kazoo.client.KazooClient.get.return_value = (b'xxx:123', None) # Need to call sync first, then put 'exit' on the queue to terminate # the loop. # # Calling run will drain event queue and populate state. app_discovery.sync() kazoo.client.KazooClient.get_children.assert_called_with( '/endpoints/appproid', watch=mock.ANY) app_discovery.exit_loop() expected = {} for (endpoint, hostport) in app_discovery.iteritems(): expected[endpoint] = hostport self.assertEqual( expected, { 'appproid.foo.1#0:tcp:http': 'xxx:123', 'appproid.foo.2#0:tcp:http': 'xxx:123' } )
def test_noexists(self): """Check that discovery establishes watch for non-existing proid.""" zkclient = kazoo.client.KazooClient() app_discovery = discovery.Discovery(zkclient, 'appproid.foo.*', 'http') app_discovery.sync() kazoo.client.KazooClient.exists.assert_called_with( '/endpoints/appproid', watch=mock.ANY)
def test_snapshot(self): """Checks that snapshot is just a copy of the internal state.""" app_discovery = discovery.Discovery(None, 'appproid.foo.*', 'http') app_discovery.state.add('foo') snapshot = app_discovery.snapshot() self.assertFalse(snapshot == app_discovery.state) self.assertEqual(set(snapshot), app_discovery.state)
def test_prefixes(self): """Dummy _prefixes() test.""" app_discovery = discovery.Discovery(None, ['appproid.foo*', 'fooproid.bar*'], '*') self.assertEqual(app_discovery._prefixes(), set(('appproid', 'fooproid')))
def test_matching_endpoints(self): """_matching_endpoints()""" app_discovery = discovery.Discovery( None, ['proid_A.kafka*', 'proid_B.zookeeper*'], '*') endpoints = [] self.assertEqual(app_discovery._matching_endpoints(endpoints), set(())) endpoints = [ discovery._join_prefix('proid_A', endpoint) for endpoint in [ 'kafka#111:tcp:http', 'kafka#111:tcp:kafka_cluster_comm', 'other#222:tcp:foo', 'asdf#333:tcp:bar', ] ] endpoints.extend([ discovery._join_prefix('proid_B', endpoint) for endpoint in [ 'zookeeper#444:tcp:zk_listener', 'zookeeper#444:tcp:http', 'xzy#555:tcp:http' ] ]) self.assertEqual( app_discovery._matching_endpoints(endpoints), set(( 'proid_A.kafka#111:tcp:http', 'proid_A.kafka#111:tcp:kafka_cluster_comm', 'proid_B.zookeeper#444:tcp:http', 'proid_B.zookeeper#444:tcp:zk_listener', )))
def test_get_endpoints_zk(self): """get_endpoints_zk(); get_endpoints()""" zkclient = mock.Mock() zkclient.get_children.return_value = [ 'foo.1#0:tcp:http', 'foo.2#0:tcp:http', 'foo.1#0:tcp:tcp', 'bar.3#0:tcp:http' ] app_discovery = discovery.Discovery( zkclient, ['proid_A.foo.1*', 'proid_B.bar.3*'], '*') self.assertEqual( app_discovery.get_endpoints_zk(), set(('proid_A.foo.1#0:tcp:http', 'proid_A.foo.1#0:tcp:tcp', 'proid_B.bar.3#0:tcp:http'))) # Test get_endpoints() def zk_get(fullpath): """Mock the zkclient.get() method.""" if fullpath.startswith( z.join_zookeeper_path(z.ENDPOINTS, 'proid_A', 'foo')): return (b'xxx:123', None) elif fullpath.startswith( z.join_zookeeper_path(z.ENDPOINTS, 'proid_B', 'bar')): return (b'yyy:987', None) else: raise ValueError(fullpath) zkclient.get = zk_get self.assertEqual(set(app_discovery.get_endpoints()), set(('xxx:123', 'xxx:123', 'yyy:987')))
def ssh(ssh, app, command): """SSH into Treadmill container.""" if ssh is None: ssh = _DEFAULT_SSH if app.find('#') == -1: # Instance is not specified, list matching and exit. raise click.BadParameter('Speficy full instance name: xxx#nnn') app_discovery = discovery.Discovery(context.GLOBAL.zk.conn, app, 'ssh') app_discovery.sync() # Restore default signal mask disabled by python spawning new thread # for Zk connection. # # TODO: should this be done as part of zkutils.connect? for sig in range(1, signal.NSIG): try: signal.signal(sig, signal.SIG_DFL) except OSError: pass # TODO: not sure how to handle mutliple instances. for (app, hostport) in app_discovery.items(): _LOGGER.info('%s :: %s', app, hostport) if hostport: host, port = hostport.split(b':') run_ssh(host, port, ssh, list(command))
def get_current_endpoints(self, is_sow=False): """This method will get the current endpoints. Based on the cell, pattern, and endpoint that the client sent in the on_message, this method will return a list of dictionaries. :param is_sow Whether this state is "SOW" (i.e. State Of the World) or not; default is False """ _LOGGER.info('Discovering endpoints in %r with pattern %s', self.cell, self.pattern) app_discovery = discovery.Discovery(self.zkclient, self.pattern, self.endpoint) endpoint_nodes = app_discovery.get_endpoints_zk() _LOGGER.debug('endpoint_nodes: %r', endpoint_nodes) endpoints = [{'name': '.'.join([app_discovery.prefix, endpoint]), 'hostport': app_discovery.resolve_endpoint(endpoint)} for endpoint in endpoint_nodes] discovered = {'sow': is_sow, 'endpoints': endpoints} return discovered
def ssh(ssh, app, command): """SSH into Treadmill container.""" if ssh is None: ssh = _DEFAULT_SSH if app.find('#') == -1: # Instance is not specified, list matching and exit. raise click.BadParameter('Specify full instance name: xxx#nnn') app_discovery = discovery.Discovery(context.GLOBAL.zk.conn, app, 'ssh') app_discovery.sync() # TODO: not sure how to handle mutliple instances. for (endpoint, hostport) in app_discovery.iteritems(): _LOGGER.info('%s :: %s', endpoint, hostport) if hostport: host, port = hostport.split(b':') run_ssh(host, port, ssh, list(command))
def vring_cmd(approot, manifest): """Run vring manager.""" context.GLOBAL.zk.conn.add_listener(zkutils.exit_on_disconnect) tm_env = appenv.AppEnvironment(approot) app = yaml.load(stream=manifest) with lc.LogContext(_LOGGER, app['name'], lc.ContainerAdapter) as log: # TODO(boysson): Remove all validation from here. utils.validate(app, [('vring', True, dict)]) ring = app['vring'] utils.validate(ring, [('rules', True, list), ('cells', True, list)]) if context.GLOBAL.cell not in ring['cells']: log.critical('cell %s not listed in vring.', context.GLOBAL.cell) sys.exit(-1) rules = ring['rules'] for rule in rules: utils.validate(rule, [('pattern', True, str), ('endpoints', True, list)]) # Create translation for endpoint name to expected port #. routing = {} for endpoint in app.get('endpoints', []): routing[endpoint['name']] = { 'port': endpoint['port'], 'proto': endpoint['proto'] } # Check that all ring endpoints are listed in the manifest. vring_endpoints = set() for rule in rules: for rule_endpoint in rule['endpoints']: if rule_endpoint not in routing: log.critical( 'vring references non-existing endpoint: [%s]', rule_endpoint) sys.exit(-1) vring_endpoints.add(rule_endpoint) # TODO: discovery is limited to one rule for now. if len(rules) != 1: log.critical('(TODO): multiple rules are not supported.') sys.exit(-1) pattern = rules[0]['pattern'] app_unique_name = appcfg.manifest_unique_name(app) app_discovery = discovery.Discovery(context.GLOBAL.zk.conn, pattern, '*') app_discovery.sync() # Restore default signal mask disabled by python spawning new # thread for Zk connection. # # TODO: should this be done as part of ZK connect? for sig in range(1, signal.NSIG): try: signal.signal(sig, signal.SIG_DFL) except RuntimeError: pass vring.run( routing, vring_endpoints, app_discovery, tm_env.rules, app['network']['vip'], app_unique_name, )
def vring_cmd(approot, manifest): """Run vring manager.""" context.GLOBAL.zk.conn.add_listener(zkutils.exit_on_disconnect) tm_env = appenv.AppEnvironment(approot) with io.open(manifest, 'r') as fd: app = json.load(fd) with lc.LogContext(_LOGGER, app['name'], lc.ContainerAdapter) as log: # TODO(boysson): Remove all validation from here. utils.validate(app, [('vring', True, dict)]) ring = app['vring'] utils.validate(ring, [('rules', True, list), ('cells', True, list)]) if context.GLOBAL.cell not in ring['cells']: log.critical('cell %s not listed in vring.', context.GLOBAL.cell) sys.exit(-1) rules = ring['rules'] for rule in rules: utils.validate(rule, [('pattern', True, str), ('endpoints', True, list)]) # Create translation for endpoint name to expected port #. routing = {} for endpoint in app.get('endpoints', []): routing[endpoint['name']] = { 'port': endpoint['port'], 'proto': endpoint['proto'] } # Check that all ring endpoints are listed in the manifest. vring_endpoints = set() for rule in rules: for rule_endpoint in rule['endpoints']: if rule_endpoint not in routing: log.critical( 'vring references non-existing endpoint: [%s]', rule_endpoint) sys.exit(-1) vring_endpoints.add(rule_endpoint) patterns = [rule['pattern'] for rule in rules] app_discovery = discovery.Discovery(context.GLOBAL.zk.conn, patterns, '*') app_discovery.sync() # Restore default signal mask disabled by python spawning new # thread for Zk connection. # # TODO: should this be done as part of ZK connect? utils.restore_signals() app_unique_name = appcfg.manifest_unique_name(app) vring.run( routing, vring_endpoints, app_discovery, tm_env.rules, app['network']['vip'], app_unique_name, )