Use fake clock for faster unit tests

databricks/sdk/clock.py

@@ -0,0 +1,49 @@
+import abc
+import time
+
+
+class Clock(metaclass=abc.ABCMeta):
+
+    @abc.abstractmethod
+    def time(self) -> float:
+        """
+        Return the current time in seconds since the Epoch.
+        Fractions of a second may be present if the system clock provides them.
+
+        :return: The current time in seconds since the Epoch.
+        """
+
+    @abc.abstractmethod
+    def sleep(self, seconds: float) -> None:
+        """
+        Return the current time in seconds since the Epoch.
+        Fractions of a second may be present if the system clock provides them.
+
+        :param seconds: The duration to sleep in seconds.
+        :return:
+        """
+
+
+class RealClock(Clock):
+    """
+    A real clock that uses the ``time`` module to get the current time and sleep.
+    """
+
+    def time(self) -> float:
+        """
+        Return the current time in seconds since the Epoch.
+        Fractions of a second may be present if the system clock provides them.
+
+        :return: The current time in seconds since the Epoch.
+        """
+        return time.time()
+
+    def sleep(self, seconds: float) -> None:
+        """
+        Return the current time in seconds since the Epoch.
+        Fractions of a second may be present if the system clock provides them.
+
+        :param seconds: The duration to sleep in seconds.
+        :return:
+        """
+        time.sleep(seconds)

databricks/sdk/config.py

@@ -11,6 +11,7 @@
 import requests
 
 from .azure import AzureEnvironment
+from .clock import Clock, RealClock
 from .credentials_provider import CredentialsProvider, DefaultCredentials
 from .environments import (ALL_ENVS, DEFAULT_ENVIRONMENT, Cloud,
                            DatabricksEnvironment)
@@ -84,13 +85,15 @@ def __init__(self,
                  credentials_provider: CredentialsProvider = None,
                  product="unknown",
                  product_version="0.0.0",
+                 clock: Clock=None,
                  **kwargs):
         self._inner = {}
         self._user_agent_other_info = []
         self._credentials_provider = credentials_provider if credentials_provider else DefaultCredentials()
         if 'databricks_environment' in kwargs:
             self.databricks_environment = kwargs['databricks_environment']
             del kwargs['databricks_environment']
+        self._clock = clock if clock is not None else RealClock()
         try:
             self._set_inner_config(kwargs)
             self._load_from_env()
@@ -317,6 +320,10 @@ def sql_http_path(self) -> Optional[str]:
         if self.warehouse_id:
             return f'/sql/1.0/warehouses/{self.warehouse_id}'
 
+    @property
+    def clock(self) -> Clock:
+        return self._clock
+
     @classmethod
     def attributes(cls) -> Iterable[ConfigAttribute]:
         """ Returns a list of Databricks SDK configuration metadata """

databricks/sdk/core.py

@@ -7,6 +7,7 @@
 
 from requests.adapters import HTTPAdapter
 
+from .clock import Clock, RealClock
 from .config import *
 # To preserve backwards compatibility (as these definitions were previously in this module)
 from .credentials_provider import *
@@ -123,7 +124,8 @@ def do(self,
             headers = {}
         headers['User-Agent'] = self._user_agent_base
         retryable = retried(timeout=timedelta(seconds=self._retry_timeout_seconds),
-                            is_retryable=self._is_retryable)
+                            is_retryable=self._is_retryable,
+                            clock=self._cfg.clock)
         return retryable(self._perform)(method,
                                         path,
                                         query=query,

databricks/sdk/retries.py

@@ -1,30 +1,34 @@
 import functools
 import logging
-import time
 from datetime import timedelta
 from random import random
 from typing import Callable, Optional, Sequence, Type
 
+from .clock import Clock, RealClock
+
 logger = logging.getLogger(__name__)
 
 
 def retried(*,
             on: Sequence[Type[BaseException]] = None,
             is_retryable: Callable[[BaseException], Optional[str]] = None,
-            timeout=timedelta(minutes=20)):
+            timeout=timedelta(minutes=20),
+            clock: Clock = None):
     has_allowlist = on is not None
     has_callback = is_retryable is not None
     if not (has_allowlist or has_callback) or (has_allowlist and has_callback):
         raise SyntaxError('either on=[Exception] or callback=lambda x: .. is required')
+    if clock is None:
+        clock = RealClock()
 
     def decorator(func):
 
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
-            deadline = time.time() + timeout.total_seconds()
+            deadline = clock.time() + timeout.total_seconds()
             attempt = 1
             last_err = None
-            while time.time() < deadline:
+            while clock.time() < deadline:
                 try:
                     return func(*args, **kwargs)
                 except Exception as err:
@@ -50,7 +54,7 @@ def wrapper(*args, **kwargs):
                         raise err
 
                     logger.debug(f'Retrying: {retry_reason} (sleeping ~{sleep}s)')
-                    time.sleep(sleep + random())
+                    clock.sleep(sleep + random())
                     attempt += 1
             raise TimeoutError(f'Timed out after {timeout}') from last_err
 

tests/clock.py

@@ -0,0 +1,16 @@
+from databricks.sdk.clock import Clock
+
+
+class FakeClock(Clock):
+    """
+    A simple clock that can be used to mock time in tests.
+    """
+
+    def __init__(self, start_time: float = 0.0):
+        self._start_time = start_time
+
+    def time(self) -> float:
+        return self._start_time
+
+    def sleep(self, seconds: float) -> None:
+        self._start_time += seconds

tests/test_core.py

@@ -26,6 +26,7 @@
 from databricks.sdk.service.iam import AccessControlRequest
 from databricks.sdk.version import __version__
 
+from .clock import FakeClock
 from .conftest import noop_credentials
 
 
@@ -422,7 +423,7 @@ def inner(h: BaseHTTPRequestHandler):
         requests.append(h.requestline)
 
     with http_fixture_server(inner) as host:
-        api_client = ApiClient(Config(host=host, token='_'))
+        api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock())
         res = api_client.do('GET', '/foo')
         assert 'foo' in res
 
@@ -445,7 +446,7 @@ def inner(h: BaseHTTPRequestHandler):
         requests.append(h.requestline)
 
     with http_fixture_server(inner) as host:
-        api_client = ApiClient(Config(host=host, token='_'))
+        api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock())
         res = api_client.do('GET', '/foo')
         assert 'foo' in res
 
@@ -462,7 +463,7 @@ def inner(h: BaseHTTPRequestHandler):
         requests.append(h.requestline)
 
     with http_fixture_server(inner) as host:
-        api_client = ApiClient(Config(host=host, token='_', retry_timeout_seconds=1))
+        api_client = ApiClient(Config(host=host, token='_', retry_timeout_seconds=1), clock=FakeClock())
         with pytest.raises(TimeoutError):
             api_client.do('GET', '/foo')
 
@@ -484,7 +485,7 @@ def inner(h: BaseHTTPRequestHandler):
         requests.append(h.requestline)
 
     with http_fixture_server(inner) as host:
-        api_client = ApiClient(Config(host=host, token='_'))
+        api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock())
         res = api_client.do('GET', '/foo')
         assert 'foo' in res
 
@@ -502,7 +503,7 @@ def inner(h: BaseHTTPRequestHandler):
         requests.append(h.requestline)
 
     with http_fixture_server(inner) as host:
-        api_client = ApiClient(Config(host=host, token='_'))
+        api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock())
         with pytest.raises(DatabricksError):
             api_client.do('GET', '/foo')
 
@@ -520,7 +521,7 @@ def inner(h: BaseHTTPRequestHandler):
         requests.append(h.requestline)
 
     with http_fixture_server(inner) as host:
-        api_client = ApiClient(Config(host=host, token='_'))
+        api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock())
         res = api_client.do('GET', '/foo')
         assert 'foo' in res