fix(oslogin): Fixing a flaky issue with OSLogin SSH tests

compute/oslogin/service_account_ssh.py

@@ -81,9 +81,9 @@ def create_ssh_key(oslogin, account, private_key_file=None, expire_time=300):
             # This method sometimes failed to work causing issues like #7277
             # Maybe retrying it with some delay will make things better
             oslogin.users().importSshPublicKey(parent=account, body=body).execute()
-        except RefreshError:
+        except RefreshError as err:
             if attempt_no == 3:
-                break
+                raise err
             time.sleep(attempt_no)
         else:
             break
@@ -126,10 +126,18 @@ def main(cmd, project, instance=None, zone=None,
     # Create a new SSH key pair and associate it with the service account.
     private_key_file = create_ssh_key(oslogin, account)
 
-    # Using the OS Login API, get the POSIX user name from the login profile
+    # Using the OS Login API, get the POSIX username from the login profile
     # for the service account.
-    profile = oslogin.users().getLoginProfile(name=account).execute()
-    username = profile.get('posixAccounts')[0].get('username')
+    for attempt_no in range(1, 4):
+        try:
+            profile = oslogin.users().getLoginProfile(name=account).execute()
+        except RefreshError as err:
+            if attempt_no == 3:
+                raise err
+            time.sleep(attempt_no)
+        else:
+            username = profile.get('posixAccounts')[0].get('username')
+            break
 
     # Create the hostname of the target instance using the instance name,
     # the zone where the instance is located, and the project that owns the

compute/oslogin/service_account_ssh_test.py

@@ -44,7 +44,7 @@ def test_main(capsys):
     project = os.environ['GOOGLE_CLOUD_PROJECT']
     test_id = 'oslogin-test-{id}'.format(id=str(random.randint(0, 1000000)))
     zone = 'us-east1-d'
-    image_family = 'projects/debian-cloud/global/images/family/debian-9'
+    image_family = 'projects/debian-cloud/global/images/family/debian-11'
     machine_type = 'zones/{zone}/machineTypes/f1-micro'.format(zone=zone)
     account_email = '{test_id}@{project}.iam.gserviceaccount.com'.format(
         test_id=test_id, project=project)
@@ -207,6 +207,9 @@ def setup_resources(
             operation=operation['name']).execute()['status'] != 'DONE':
         time.sleep(5)
 
+    # Wait for the OS of the instance to be ready to accept SSH connections
+    time.sleep(10)
+
     # Grant the service account osLogin access on the test instance.
     compute.instances().setIamPolicy(
         project=project,