From 82f99687192e7dc14530a8a604cdb2f1ac4208f4 Mon Sep 17 00:00:00 2001
From: Cory Johns <johnsca@gmail.com>
Date: Mon, 17 Sep 2018 12:51:37 -0400
Subject: [PATCH] Retry ssh in manual provision test (#265)

Also fix reference-before-assignment error on failure
---
 juju/provisioner.py             |  1 +
 tests/integration/test_model.py | 32 +++++++++++++++++++++++---------
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/juju/provisioner.py b/juju/provisioner.py
index 91747a453..1be07bfbb 100644
--- a/juju/provisioner.py
+++ b/juju/provisioner.py
@@ -158,6 +158,7 @@ def _init_ubuntu_user(self):
         # TODO: Test this on an image without the ubuntu user setup.
 
         auth_user = self.user
+        ssh = None
         try:
             # Run w/o allocating a pty, so we fail if sudo prompts for a passwd
             ssh = self._get_ssh_client(
diff --git a/tests/integration/test_model.py b/tests/integration/test_model.py
index 675637eee..205a77a4c 100644
--- a/tests/integration/test_model.py
+++ b/tests/integration/test_model.py
@@ -2,6 +2,7 @@
 import mock
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
+import paramiko
 
 from juju.client.client import ConfigValue, ApplicationFacade
 from juju.model import Model, ModelObserver
@@ -167,7 +168,9 @@ async def test_add_manual_machine_ssh(event_loop):
 
         profile = client.profiles.create(
             test_name,
-            config={'user.user-data': '#cloud-config\nssh_authorized_keys:\n- {}'.format(public_key)},
+            config={'user.user-data': '#cloud-config\n'
+                                      'ssh_authorized_keys:\n'
+                                      '- {}'.format(public_key)},
             devices={
                 'root': {'path': '/', 'pool': 'default', 'type': 'disk'},
                 'eth0': {
@@ -206,19 +209,28 @@ def wait_for_network(container, timeout=30):
             return None
 
         host = wait_for_network(container)
+        assert host, 'Failed to get address for machine'
 
         # HACK: We need to give sshd a chance to bind to the interface,
         # and pylxd's container.execute seems to be broken and fails and/or
         # hangs trying to properly check if the service is up.
         time.sleep(5)
 
-        if host:
-            # add a new manual machine
-            machine1 = await model.add_machine(spec='ssh:{}@{}:{}'.format(
-                "ubuntu",
-                host['address'],
-                private_key_path,
-            ))
+        for attempt in range(1, 4):
+            try:
+                # add a new manual machine
+                machine1 = await model.add_machine(spec='ssh:{}@{}:{}'.format(
+                    "ubuntu",
+                    host['address'],
+                    private_key_path,
+                ))
+            except paramiko.ssh_exception.NoValidConnectionsError:
+                if attempt == 3:
+                    raise
+                # retry the ssh connection a few times if it fails
+                time.sleep(attempt * 5)
+            else:
+                break
 
             assert len(model.machines) == 1
 
@@ -307,7 +319,9 @@ async def test_explicit_loop_threaded(event_loop):
         with ThreadPoolExecutor(1) as executor:
             f = executor.submit(
                 new_loop.run_until_complete,
-                _deploy_in_loop(new_loop, model_name, model._connector.jujudata))
+                _deploy_in_loop(new_loop,
+                                model_name,
+                                model._connector.jujudata))
             f.result()
         await model._wait_for_new('application', 'ubuntu')
         assert 'ubuntu' in model.applications