chore: refactor perform_refresh and add tests (#144)

* chore(tests): refactor perform_refresh and add tests

* chore: move stubs to outside of tests

* use run_coroutine_threadsafe to call _perform_refresh in tests

* linting

Author: shubha-rajan

google/cloud/sql/connector/instance_connection_manager.py

@@ -383,35 +383,32 @@ async def _perform_refresh(self) -> asyncio.Task:
 
         refresh_task = self._loop.create_task(self._get_instance_data())
 
-        def _refresh_callback(task: asyncio.Task) -> None:
+        try:
+            await refresh_task
+        except Exception as e:
+            logger.exception(
+                "An error occurred while performing refresh. Retrying in 60s.",
+                exc_info=e,
+            )
+            instance_data = None
             try:
-                task.result()
-            except Exception as e:
-                logger.exception(
-                    "An error occurred while performing refresh. Retrying in 60s.",
-                    exc_info=e,
-                )
-                instance_data = None
-                try:
-                    instance_data = self._current.result()
-                except Exception:
-                    # Current result is invalid, no-op
-                    logger.debug("Current instance data is invalid.")
-                if (
-                    instance_data is None
-                    or instance_data.expiration < datetime.datetime.now()
-                ):
-                    self._current = task
-                    # TODO: Implement force refresh method and a rate-limiter for perform_refresh
-                    # Retry by scheduling a refresh 60s from now.
-                    self._next = self._loop.create_task(self._schedule_refresh(60))
-
-            else:
+                instance_data = await self._current
+            except Exception:
+                # Current result is invalid, no-op
+                logger.debug("Current instance data is invalid.")
+            if (
+                instance_data is None
+                or instance_data.expiration < datetime.datetime.now()
+            ):
                 self._current = refresh_task
-                # Ephemeral certificate expires in 1 hour, so we schedule a refresh to happen in 55 minutes.
-                self._next = self._loop.create_task(self._schedule_refresh())
+                # TODO: Implement force refresh method and a rate-limiter for perform_refresh
+                # Retry by scheduling a refresh 60s from now.
+                self._next = self._loop.create_task(self._schedule_refresh(60))
 
-        refresh_task.add_done_callback(_refresh_callback)
+        else:
+            self._current = refresh_task
+            # Ephemeral certificate expires in 1 hour, so we schedule a refresh to happen in 55 minutes.
+            self._next = self._loop.create_task(self._schedule_refresh())
 
         return refresh_task
 

noxfile.py

@@ -65,6 +65,7 @@ def default(session, path):
         "py.test",
         # "--cov=util",
         # "--cov=connector",
+        "-v",
         "--cov-append",
         "--cov-config=.coveragerc",
         "--cov-report=",

tests/unit/test_instance_connection_manager.py

@@ -15,6 +15,8 @@
 """
 
 import asyncio
+import datetime
+from typing import Any
 import pytest  # noqa F401 Needed to run the tests
 from google.cloud.sql.connector.instance_connection_manager import (
     InstanceConnectionManager,
@@ -32,6 +34,19 @@ def icm(
     return icm
 
 
+class MockMetadata:
+    def __init__(self, expiration: datetime.datetime) -> None:
+        self.expiration = expiration
+
+
+async def _get_metadata_success(*args: Any, **kwargs: Any) -> MockMetadata:
+    return MockMetadata(datetime.datetime.now() + datetime.timedelta(minutes=10))
+
+
+async def _get_metadata_error(*args: Any, **kwargs: Any) -> None:
+    raise Exception("something went wrong...")
+
+
 def test_InstanceConnectionManager_init(async_loop: asyncio.AbstractEventLoop) -> None:
     """
     Test to check whether the __init__ method of InstanceConnectionManager
@@ -53,13 +68,67 @@ def test_InstanceConnectionManager_init(async_loop: asyncio.AbstractEventLoop) -
 
 
 @pytest.mark.asyncio
-async def test_InstanceConnectionManager_perform_refresh(
-    icm: InstanceConnectionManager, async_loop: asyncio.AbstractEventLoop
+async def test_perform_refresh_replaces_result(icm: InstanceConnectionManager) -> None:
+    """
+    Test to check whether _perform_refresh replaces a valid result with another valid result
+    """
+
+    # stub _get_instance_data to return a "valid" MockMetadata object
+    setattr(icm, "_get_instance_data", _get_metadata_success)
+    new_task = asyncio.run_coroutine_threadsafe(
+        icm._perform_refresh(), icm._loop
+    ).result(timeout=10)
+
+    assert icm._current == new_task
+    assert isinstance(icm._current.result(), MockMetadata)
+
+
+@pytest.mark.asyncio
+async def test_perform_refresh_wont_replace_valid_result_with_invalid(
+    icm: InstanceConnectionManager,
 ) -> None:
     """
-    Test to check whether _perform_refresh works as described given valid
-    conditions.
+    Test to check whether _perform_refresh won't replace a valid _current
+    value with an invalid one
     """
-    task = await icm._perform_refresh()
 
-    assert isinstance(task, asyncio.Task)
+    # stub _get_instance_data to return a "valid" MockMetadata object
+    setattr(icm, "_get_instance_data", _get_metadata_success)
+    icm._current = asyncio.run_coroutine_threadsafe(
+        icm._perform_refresh(), icm._loop
+    ).result(timeout=10)
+    old_task = icm._current
+
+    # stub _get_instance_data to throw an error, then await _perform_refresh
+    setattr(icm, "_get_instance_data", _get_metadata_error)
+    asyncio.run_coroutine_threadsafe(icm._perform_refresh(), icm._loop).result(
+        timeout=10
+    )
+
+    assert icm._current == old_task
+    assert isinstance(icm._current.result(), MockMetadata)
+
+
+@pytest.mark.asyncio
+async def test_perform_refresh_replaces_invalid_result(
+    icm: InstanceConnectionManager,
+) -> None:
+    """
+    Test to check whether _perform_refresh will replace an invalid refresh result with
+    a valid one
+    """
+
+    # stub _get_instance_data to throw an error
+    setattr(icm, "_get_instance_data", _get_metadata_error)
+    icm._current = asyncio.run_coroutine_threadsafe(
+        icm._perform_refresh(), icm._loop
+    ).result(timeout=10)
+
+    # stub _get_instance_data to return a MockMetadata instance
+    setattr(icm, "_get_instance_data", _get_metadata_success)
+    new_task = asyncio.run_coroutine_threadsafe(
+        icm._perform_refresh(), icm._loop
+    ).result(timeout=10)
+
+    assert icm._current == new_task
+    assert isinstance(icm._current.result(), MockMetadata)