From c5dec1285ce9655b0253b9ce459b0b3aa92dd797 Mon Sep 17 00:00:00 2001
From: Michael Thornton <mbthornton@lbl.gov>
Date: Tue, 17 Dec 2024 11:39:18 -0800
Subject: [PATCH] Add unit test for get_token retry behavior

---
 nmdc_automation/api/nmdcapi.py |  6 +++---
 tests/test_nmdcapi.py          | 24 ++++++++++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/nmdc_automation/api/nmdcapi.py b/nmdc_automation/api/nmdcapi.py
index f5560bd6..6347e2a2 100755
--- a/nmdc_automation/api/nmdcapi.py
+++ b/nmdc_automation/api/nmdcapi.py
@@ -75,14 +75,14 @@ def _get_token(self, *args, **kwargs):
         return _get_token
 
     @retry(
-        wait=wait_exponential(multiplier=1, min=30, max=240),
-        stop=stop_after_attempt(4),
+        wait=wait_exponential(multiplier=4, min=8, max=120),
+        stop=stop_after_attempt(6),
         reraise=True,
     )
     def get_token(self):
         """
         Get a token using a client id/secret.
-        Retries up to 4 times with exponential backoff.
+        Retries up to 6 times with exponential backoff.
         """
         h = {
             "accept": "application/json",
diff --git a/tests/test_nmdcapi.py b/tests/test_nmdcapi.py
index 72f2f05f..4bd70de5 100644
--- a/tests/test_nmdcapi.py
+++ b/tests/test_nmdcapi.py
@@ -86,3 +86,27 @@ def test_jobs(requests_mock, site_config_file, mock_api):
     requests_mock.post(url, json={}, status_code=409)
     resp = n.claim_job("abc")
     assert resp["claimed"] is True
+
+
+def test_nmdcapi_get_token_with_retry(requests_mock, site_config_file):
+    n = nmdcapi(site_config_file)
+    token_url = "http://localhost:8000/token"
+
+    requests_mock.post(
+        token_url, [{"status_code": 401, "json": {"error": "Unauthorized"}},
+                    {"status_code": 200, "json": {
+                        "access_token": "mocked_access_token",
+                        "expires": {"days": 1},
+                    }}]
+    )
+    # sanity check
+    assert n.token is None
+    assert n.expires_at == 0
+
+    # Call method under test - should retry and succeed
+    n.get_token()
+
+    # Check that the token was set
+    assert n.token == "mocked_access_token"
+    assert n.expires_at > 0
+