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,
)
