From a86d380568d3ad96585163c8e7134188f4e3979f Mon Sep 17 00:00:00 2001
From: Dan Lee <71398022+dandhlee@users.noreply.github.com>
Date: Tue, 16 Aug 2022 08:50:56 -0500
Subject: [PATCH] fix: use backoff exception to gracefully exit and fail
 (#8257)

* fix: use backoff exception to gracefully exit and fail

* fix: using specific tuple of exceptions to track

* fix: update syntax

* fix: lint

* fix: lint2

* fix: update lint

* chore(deps): pin dependency for versions

Co-authored-by: Maciej Strzelczyk <strzelczyk@google.com>
---
 compute/oslogin/requirements-test.txt       |  3 ++-
 compute/oslogin/service_account_ssh_test.py | 21 +++++++++++----------
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/compute/oslogin/requirements-test.txt b/compute/oslogin/requirements-test.txt
index 8688f74ea0b1..166cd92b46be 100644
--- a/compute/oslogin/requirements-test.txt
+++ b/compute/oslogin/requirements-test.txt
@@ -1,3 +1,4 @@
+backoff===1.11.1; python_version < "3.7"
+backoff===2.0.0; python_version >= "3.7"
 pytest==7.0.1
 mock==4.0.3
-retrying==1.3.3
diff --git a/compute/oslogin/service_account_ssh_test.py b/compute/oslogin/service_account_ssh_test.py
index 46ffb52a188e..9a862a682eb0 100644
--- a/compute/oslogin/service_account_ssh_test.py
+++ b/compute/oslogin/service_account_ssh_test.py
@@ -16,11 +16,13 @@
 import json
 import os
 import random
+from subprocess import CalledProcessError
 import time
 
+import backoff
+from google.auth.exceptions import RefreshError
 from google.oauth2 import service_account
 import googleapiclient.discovery
-from retrying import retry
 
 from service_account_ssh import main
 
@@ -84,8 +86,12 @@ def test_main(capsys):
         'oslogin', 'v1', cache_discovery=False, credentials=credentials)
     account = 'users/' + account_email
 
-    @retry(wait_exponential_multiplier=1000, wait_exponential_max=300000,
-           stop_max_attempt_number=10)
+    # More exceptions could be raised, keeping track of ones I could
+    # find for now.
+    @backoff.on_exception(backoff.expo,
+                          (CalledProcessError,
+                           RefreshError),
+                          max_tries=5)
     def ssh_login():
         main(cmd, project, test_id, zone, oslogin, account, hostname)
         out, _ = capsys.readouterr()
@@ -93,13 +99,8 @@ def ssh_login():
         assert assert_value in out
 
     # Test SSH to the instance.
-    try:
-        ssh_login()
-    except Exception:
-        raise Exception('SSH to the test instance failed.')
-
-    finally:
-        cleanup_resources(compute, iam, project, test_id, zone, account_email)
+    ssh_login()
+    cleanup_resources(compute, iam, project, test_id, zone, account_email)
 
 
 def setup_resources(