315
323
mock_sleep.assert_called_with(1)
326
class TestGetAdminSecret(unittest.TestCase):
328
def mock_environment_file(self, contents, juju_env=None):
329
"""Create a mock environment file containing the given contents."""
330
# Create a temporary home that will be used by get_admin_secret().
331
home = tempfile.mkdtemp()
332
self.addCleanup(shutil.rmtree, home)
333
# Create a juju home.
334
juju_home = os.path.join(home, '.juju')
336
# Set up an environments file with the given contents.
337
environments_path = os.path.join(juju_home, 'environments.yaml')
338
with open(environments_path, 'w') as environments_file:
339
environments_file.write(contents)
340
# Return a mock object patching the environment context with the
341
# temporary HOME and JUJU_ENV.
342
environ = {'HOME': home}
343
if juju_env is not None:
344
environ['JUJU_ENV'] = juju_env
345
# The returned object can be used as a context manager.
346
return mock.patch('os.environ', environ)
349
def assert_error(self, error):
350
"""Ensure a ValueError is raised in the context block.
352
Also check that the exception includes the expected error message.
354
with self.assertRaises(ValueError) as context_manager:
356
self.assertIn(error, str(context_manager.exception))
358
def test_no_env_name(self):
359
# A ValueError is raised if the env name is not included in the
360
# environment context.
361
expected = 'Unable to retrieve the current environment name.'
362
with self.mock_environment_file(''):
363
with self.assert_error(expected):
366
def test_no_file(self):
367
# A ValueError is raised if the environments file is not found.
368
home = tempfile.mkdtemp()
369
self.addCleanup(shutil.rmtree, home)
370
expected = 'Unable to open environments file: [Errno 2] No such file'
371
with mock.patch('os.environ', {'HOME': home, 'JUJU_ENV': 'ec2'}):
372
with self.assert_error(expected):
375
def test_invalid_yaml(self):
376
# A ValueError is raised if the environments file is not well formed.
377
with self.mock_environment_file(':', juju_env='ec2'):
378
with self.assert_error('Unable to parse environments file:'):
381
def test_invalid_yaml_contents(self):
382
# A ValueError is raised if the environments file is not well formed.
383
with self.mock_environment_file('a-string', juju_env='ec2'):
384
with self.assert_error('Invalid YAML contents: a-string'):
387
def test_no_env(self):
388
# A ValueError is raised if the environment is not found in the YAML.
389
contents = yaml.safe_dump({'environments': {'local': {}}})
390
with self.mock_environment_file(contents, juju_env='ec2'):
391
with self.assert_error('Environment ec2 not found'):
394
def test_no_admin_secret(self):
395
# A ValueError is raised if the admin secret is not included in the
397
contents = yaml.safe_dump({'environments': {'ec2': {}}})
398
with self.mock_environment_file(contents, juju_env='ec2'):
399
with self.assert_error('Admin secret not found'):
403
# The environment admin secret is correctly returned.
404
contents = yaml.safe_dump({
405
'environments': {'ec2': {'admin-secret': 'Secret!'}},
407
with self.mock_environment_file(contents, juju_env='ec2'):
408
self.assertEqual('Secret!', get_admin_secret())
411
@mock.patch('helpers.ssh')
412
class TestStopServices(unittest.TestCase):
414
def test_single_service(self, mock_ssh):
415
# An ssh command is executed to stop a single service.
416
stop_services('my-host-name', ['foo'])
417
mock_ssh.assert_called_once_with(
418
'ubuntu@my-host-name', 'sudo', 'service', 'foo', 'stop')
420
def test_multiple_services(self, mock_ssh):
421
# An ssh command is executed for each given service.
422
stop_services('127.0.0.1', ['foo', 'bar'])
423
self.assertEqual(2, mock_ssh.call_count)
425
mock.call('ubuntu@127.0.0.1', 'sudo', 'service', 'foo', 'stop'),
426
mock.call('ubuntu@127.0.0.1', 'sudo', 'service', 'bar', 'stop'),
428
mock_ssh.assert_has_calls(expected)
318
431
@mock.patch('helpers.juju_status')
319
class TestWaitForService(unittest.TestCase):
432
class TestWaitForUnit(unittest.TestCase):
321
434
address = 'unit.example.com'
322
435
service = 'test-service'
395
508
unit_info = wait_for_unit(self.service)
396
509
self.assertEqual(self.address, unit_info['public-address'])
397
510
self.assertEqual(1, mock_juju_status.call_count)
513
@mock.patch('helpers.websocket.create_connection')
514
class TestWebSocketClient(unittest.TestCase):
516
url = 'wss://example.com:17070'
519
self.client = WebSocketClient(self.url)
521
def test_connect(self, mock_create_connection):
522
# The WebSocket connection is correctly established.
523
self.client.connect()
524
mock_create_connection.assert_called_once_with(self.url)
526
def test_send(self, mock_create_connection):
527
# A request is correctly sent, and a response is returned.
528
request = {'request': 'foo'}
529
expected_request = json.dumps(request)
530
expected_response = {'response': 'bar'}
531
mock_connection = mock_create_connection()
532
mock_connection.recv.return_value = json.dumps(expected_response)
533
self.client.connect()
534
response = self.client.send(request)
535
self.assertEqual(expected_response, response)
536
mock_connection.send.assert_called_once_with(expected_request)
538
def test_close(self, mock_create_connection):
539
# A WebSocket connection is correctly terminated.
540
self.client.connect()
542
mock_create_connection().close.assert_called_once_with()