Add oauth token caching to DNAAS client

src/ralph/dns/dnsaas.py

@@ -3,6 +3,7 @@
 import logging
 from datetime import datetime, timedelta
 from functools import wraps
+from threading import Lock
 from typing import List, Optional, Tuple, Union
 from urllib.parse import parse_qs, urlencode, urljoin, urlsplit
 
@@ -41,6 +42,9 @@ def wrapper(self, *args, **kwargs):
 
 
 class DNSaaS:
+    _oauth_tokens = {}
+    _oauth_token_lock = Lock()
+
     def __init__(self, headers: dict = None):
         self.session = requests.Session()
         _headers = {
@@ -52,25 +56,39 @@ def __init__(self, headers: dict = None):
             _headers.update(headers)
         self.session.headers.update(_headers)
 
-    def _get_oauth_token(self):
+    def _get_oauth_token(self, force_refresh=False):
         client_id = settings.OAUTH_CLIENT_ID
         secret = settings.OAUTH_SECRET
         token_url = settings.OAUTH_TOKEN_URL
-        client = BackendApplicationClient(client_id=client_id)
-        oauth = OAuth2Session(client=client)
-        try:
-            token = oauth.fetch_token(
-                token_url=token_url, client_id=client_id, client_secret=secret
-            )
-        except CustomOAuth2Error as e:
-            logger.error(str(e))
+        cache_key = (token_url, client_id)
+
+        with self._oauth_token_lock:
+            now = datetime.now()
+            cached_token = self._oauth_tokens.get(cache_key)
+            if not force_refresh and cached_token and now < cached_token["expiration"]:
+                self.token_expiration = cached_token["expiration"]
+                return cached_token["access_token"]
 
-        expire_in = token.get("expires_in")
-        self.token_expiration = datetime.now() + timedelta(0, expire_in - 60)
-        return token.get("access_token")
+            client = BackendApplicationClient(client_id=client_id)
+            oauth = OAuth2Session(client=client)
+            try:
+                token = oauth.fetch_token(
+                    token_url=token_url, client_id=client_id, client_secret=secret
+                )
+            except CustomOAuth2Error:
+                logger.exception("Fetching an OAuth token for DNSaaS failed")
+                raise
+
+            expires_in = token.get("expires_in")
+            self.token_expiration = now + timedelta(seconds=max(expires_in - 60, 0))
+            self._oauth_tokens[cache_key] = {
+                "access_token": token.get("access_token"),
+                "expiration": self.token_expiration,
+            }
+            return token.get("access_token")
 
     def _update_oauth_token(self):
-        token = self._get_oauth_token()
+        token = self._get_oauth_token(force_refresh=True)
         self.session.headers["Authorization"] = "Bearer {}".format(token)
 
     def _verify_oauth_token_validity(self):

src/ralph/dns/tests.py

@@ -1,8 +1,10 @@
 # -*- coding: utf-8 -*-
+from datetime import datetime, timedelta
 from unittest.mock import patch
 
 from django.db import transaction
 from django.test import override_settings, TestCase, TransactionTestCase
+from oauthlib.oauth2.rfc6749.errors import CustomOAuth2Error
 
 from ralph.assets.tests.factories import ConfigurationClassFactory, EthernetFactory
 from ralph.data_center.models import BaseObjectCluster, DataCenterAsset
@@ -15,6 +17,63 @@
 from ralph.virtual.tests.factories import VirtualServerFactory
 
 
+@override_settings(
+    OAUTH_CLIENT_ID="ralph",
+    OAUTH_SECRET="secret",
+    OAUTH_TOKEN_URL="https://oauth.example.com/token",
+)
+class TestOAuthToken(TestCase):
+    def setUp(self):
+        DNSaaS._oauth_tokens.clear()
+
+    def tearDown(self):
+        DNSaaS._oauth_tokens.clear()
+
+    @patch("ralph.dns.dnsaas.OAuth2Session")
+    def test_reuses_token_between_clients_until_it_expires(self, oauth_session_mock):
+        oauth_session_mock.return_value.fetch_token.return_value = {
+            "access_token": "token",
+            "expires_in": 3600,
+        }
+
+        first_client = DNSaaS()
+        second_client = DNSaaS()
+
+        oauth_session_mock.return_value.fetch_token.assert_called_once()
+        self.assertEqual(first_client.session.headers["Authorization"], "Bearer token")
+        self.assertEqual(second_client.session.headers["Authorization"], "Bearer token")
+
+    @patch("ralph.dns.dnsaas.OAuth2Session")
+    def test_fetches_new_token_when_cached_token_has_expired(self, oauth_session_mock):
+        oauth_session_mock.return_value.fetch_token.side_effect = [
+            {"access_token": "old-token", "expires_in": 3600},
+            {"access_token": "new-token", "expires_in": 3600},
+        ]
+        first_client = DNSaaS()
+        cached_token = next(iter(DNSaaS._oauth_tokens.values()))
+        cached_token["expiration"] = datetime.now() - timedelta(seconds=1)
+
+        second_client = DNSaaS()
+
+        self.assertEqual(
+            first_client.session.headers["Authorization"], "Bearer old-token"
+        )
+        self.assertEqual(
+            second_client.session.headers["Authorization"], "Bearer new-token"
+        )
+        self.assertEqual(oauth_session_mock.return_value.fetch_token.call_count, 2)
+
+    @patch("ralph.dns.dnsaas.OAuth2Session")
+    def test_preserves_oauth_error_when_fetching_token_fails(self, oauth_session_mock):
+        oauth_error = CustomOAuth2Error(description="rate limited")
+        oauth_session_mock.return_value.fetch_token.side_effect = oauth_error
+
+        with self.assertRaises(CustomOAuth2Error) as raised:
+            DNSaaS()
+
+        self.assertIs(raised.exception, oauth_error)
+
+
 class TestGetDnsRecords(TestCase):
     @patch.object(DNSaaS, "_get_oauth_token")
     def setUp(self, mocked):