Return refresh response instead of throwing exception

Author: asloobq

examples/sample_bidstream_client.py

@@ -21,7 +21,11 @@ def _usage():
 ad_token = sys.argv[5]
 
 client = BidstreamClient(base_url, auth_key, secret_key)
-client.refresh()
+refresh_response = client.refresh()
+if not refresh_response.success:
+    print('Failed to refresh keys due to =', refresh_response.reason)
+    sys.exit(1)
+
 decrypt_result = client.decrypt_token_into_raw_uid(ad_token, domain_name)
 
 print('Status =', decrypt_result.status)

examples/sample_sharing_client.py

@@ -20,7 +20,11 @@ def _usage():
 raw_uid = sys.argv[4]
 
 client = SharingClient(base_url, auth_key, secret_key)
-client.refresh()
+refresh_response = client.refresh()
+if not refresh_response.success:
+    print('Failed to refresh keys due to =', refresh_response.reason)
+    sys.exit(1)
+
 encryption_data_response = client.encrypt_raw_uid_into_token(raw_uid)
 
 print('New Sharing Token =', encryption_data_response.encrypted_data)

tests/test_bidstream_client.py

@@ -3,6 +3,11 @@
 
 from test_utils import *
 from uid2_client import BidstreamClient, EncryptionError, Uid2Base64UrlCoder, DecryptionStatus
+from uid2_client.refresh_response import RefreshResponse
+
+
+def create_default_refresh_keys_success():
+    return RefreshResponse.make_success(create_default_key_collection([master_key, site_key]))
 
 
 @patch('uid2_client.bid_stream_client.refresh_bidstream_keys')
@@ -34,16 +39,17 @@ def test_smoke_test(self, mock_refresh_bidstream_keys):  # SmokeTest
         for expected_scope, expected_version in test_cases_all_scopes_all_versions:
             with self.subTest(expected_scope=expected_scope, expected_version=expected_version):
                 token = generate_uid_token(expected_scope, expected_version)
-                mock_refresh_bidstream_keys.return_value = create_key_collection(expected_scope)
+                mock_refresh_bidstream_keys.return_value = RefreshResponse.make_success(
+                    create_key_collection(expected_scope))
                 self._client.refresh()
                 self.decrypt_and_assert_success(token, expected_version, expected_scope)
 
     def test_phone_uids(self, mock_refresh_bidstream_keys):  # PhoneTest
         for expected_scope, expected_version in test_cases_all_scopes_v3_v4_versions:
             with self.subTest(expected_scope=expected_scope, expected_version=expected_version):
-                mock_refresh_bidstream_keys.return_value = EncryptionKeysCollection([master_key, site_key],
+                mock_refresh_bidstream_keys.return_value = RefreshResponse.make_success(EncryptionKeysCollection([master_key, site_key],
                                                                                     expected_scope, site_id, 1,
-                                                                                  99999, 86400)
+                                                                                  99999, 86400))
                 self._client.refresh()
                 token = generate_uid_token(expected_scope, expected_version, phone_uid)
                 self.assertEqual(IdentityType.Phone, get_identity_type(token))
@@ -60,10 +66,10 @@ def test_token_lifetime_too_long_for_bidstream(self, mock_refresh_bidstream_keys
         for expected_scope, expected_version in test_cases_all_scopes_all_versions:
             with self.subTest(expected_scope=expected_scope, expected_version=expected_version):
                 token = generate_uid_token(expected_scope, expected_version, expires_at=expires_in_sec)
-                mock_refresh_bidstream_keys.return_value = EncryptionKeysCollection([master_key, site_key],
+                mock_refresh_bidstream_keys.return_value = RefreshResponse.make_success(EncryptionKeysCollection([master_key, site_key],
                                                                                   expected_scope, site_id, 1,
                                                                                   99999, 86400,
-                                                                                  max_bidstream_lifetime_seconds=max_bidstream_lifetime)
+                                                                                  max_bidstream_lifetime_seconds=max_bidstream_lifetime))
                 self._client.refresh()
                 result = self._client.decrypt_token_into_raw_uid(token, None)
                 self.assertFalse(result.success)
@@ -74,9 +80,8 @@ def test_token_generated_in_the_future_to_simulate_clock_skew(self, mock_refresh
         for expected_scope, expected_version in test_cases_all_scopes_all_versions:
             with self.subTest(expected_scope=expected_scope, expected_version=expected_version):
                 token = generate_uid_token(expected_scope, expected_version, created_at=created_at_future)
-                mock_refresh_bidstream_keys.return_value = EncryptionKeysCollection([master_key, site_key],
-                                                                                  expected_scope, site_id, 1,
-                                                99999, 86400)
+                mock_refresh_bidstream_keys.return_value = RefreshResponse.make_success(EncryptionKeysCollection(
+                    [master_key, site_key], expected_scope, site_id, 1,99999, 86400))
                 self._client.refresh()
                 result = self._client.decrypt_token_into_raw_uid(token, None)
                 self.assertFalse(result.success)
@@ -87,9 +92,9 @@ def test_token_generated_in_the_future_within_allowed_clock_skew(self, mock_refr
         for expected_scope, expected_version in test_cases_all_scopes_all_versions:
             with self.subTest(expected_scope=expected_scope, expected_version=expected_version):
                 token = generate_uid_token(expected_scope, expected_version, expires_at=created_at_future)
-                mock_refresh_bidstream_keys.return_value = EncryptionKeysCollection([master_key, site_key],
-                                                                                  expected_scope, site_id, 1,
-                                                99999, 86400)
+                mock_refresh_bidstream_keys.return_value = RefreshResponse.make_success(
+                    EncryptionKeysCollection([master_key, site_key], expected_scope, site_id, 1,
+                                                99999, 86400))
                 self._client.refresh()
                 result = self._client.decrypt_token_into_raw_uid(token, None)
                 self.assertIsNotNone(result)
@@ -98,7 +103,7 @@ def test_token_generated_in_the_future_within_allowed_clock_skew(self, mock_refr
 
     def test_empty_keys(self, mock_refresh_bidstream_keys):  # EmptyKeyContainer
         token = generate_uid_token(IdentityScope.UID2, AdvertisingTokenVersion.ADVERTISING_TOKEN_V3)
-        mock_refresh_bidstream_keys.return_value = None
+        mock_refresh_bidstream_keys.return_value = RefreshResponse.make_error('Exception')
         self._client.refresh()
         result = self._client.decrypt_token_into_raw_uid(token, None)
         self.assertFalse(result.success)
@@ -110,7 +115,7 @@ def get_post_refresh_keys_response_with_key_expired():
                                         keyset_id=99999)
             site_key_expired = EncryptionKey(site_key_id, site_id, created=now, activates=now - dt.timedelta(hours=2), expires=now - dt.timedelta(hours=1), secret=site_secret,
                                         keyset_id=99999)
-            return create_default_key_collection([master_key_expired, site_key_expired])
+            return RefreshResponse.make_success(create_default_key_collection([master_key_expired, site_key_expired]))
 
         mock_refresh_bidstream_keys.return_value = get_post_refresh_keys_response_with_key_expired()
         self._client.refresh()
@@ -123,7 +128,7 @@ def test_not_authorized_for_master_key(self, mock_refresh_bidstream_keys):  #Not
         def get_post_refresh_keys_response_with_key_expired():
             another_master_key = EncryptionKey(master_key_id + site_key_id + 1, -1, created=now, activates=now, expires=now + dt.timedelta(hours=1), secret=master_secret)
             another_site_key = EncryptionKey(master_key_id + site_key_id + 2, site_id, created=now, activates=now, expires=now + dt.timedelta(hours=1), secret=site_secret)
-            return create_default_key_collection([another_master_key, another_site_key])
+            return RefreshResponse.make_success(create_default_key_collection([another_master_key, another_site_key]))
 
         mock_refresh_bidstream_keys.return_value = get_post_refresh_keys_response_with_key_expired()
         self._client.refresh()
@@ -134,7 +139,7 @@ def get_post_refresh_keys_response_with_key_expired():
         self.assertEqual(result.status, DecryptionStatus.NOT_AUTHORIZED_FOR_MASTER_KEY)
 
     def test_invalid_payload(self, mock_refresh_bidstream_keys):  #InvalidPayload
-        mock_refresh_bidstream_keys.return_value = create_default_key_collection([master_key, site_key])
+        mock_refresh_bidstream_keys.return_value = create_default_refresh_keys_success()
         self._client.refresh()
         token = generate_uid_token(IdentityScope.UID2, AdvertisingTokenVersion.ADVERTISING_TOKEN_V4)
         payload = Uid2Base64UrlCoder.decode(token)
@@ -145,7 +150,7 @@ def test_invalid_payload(self, mock_refresh_bidstream_keys):  #InvalidPayload
         self.assertEqual(result.status, DecryptionStatus.INVALID_PAYLOAD)
 
     def test_token_expiry_custom_decryption_time(self, mock_refresh_bidstream_keys):  #TokenExpiryAndCustomNow
-        mock_refresh_bidstream_keys.return_value = create_default_key_collection([master_key, site_key])
+        mock_refresh_bidstream_keys.return_value = create_default_refresh_keys_success()
         self._client.refresh()
 
         expires_at = now - dt.timedelta(days=60)
@@ -163,7 +168,7 @@ def test_token_expiry_custom_decryption_time(self, mock_refresh_bidstream_keys):
         self.assertEqual(result.advertising_token_version, AdvertisingTokenVersion.ADVERTISING_TOKEN_V4)
 
     def test_refresh_keys(self, mock_refresh_bidstream_keys):
-        mock_refresh_bidstream_keys.return_value = create_default_key_collection([master_key])
+        mock_refresh_bidstream_keys.return_value = RefreshResponse.make_success(create_default_key_collection([master_key]))
         self._client.refresh()
         mock_refresh_bidstream_keys.assert_called_once_with(self._CONST_BASE_URL, self._CONST_API_KEY,
                                                             client_secret_bytes)

tests/test_client.py

@@ -4,6 +4,7 @@
 from test_utils import *
 from uid2_client import *
 from uid2_client.euid_client_factory import EuidClientFactory
+from uid2_client.refresh_response import RefreshResponse
 from uid2_client.uid2_client_factory import Uid2ClientFactory
 
 
@@ -19,12 +20,19 @@ def setUp(self):
                                                         99999, 86400)
 
     def test_refresh(self, mock_refresh_keys_util):
-        mock_refresh_keys_util.return_value = self._key_collection
+        mock_refresh_keys_util.return_value = RefreshResponse.make_success(self._key_collection)
         client = Uid2ClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)
         client.refresh_keys()
         mock_refresh_keys_util.assert_called_once_with(self._CONST_BASE_URL, self._CONST_API_KEY, base64.b64decode(client_secret))
         self.assertEqual(client._keys, self._key_collection)
 
+    def test_refresh_fail(self, mock_refresh_keys_util):
+        mock_refresh_keys_util.return_value = RefreshResponse.make_error('Exception msg')
+        client = Uid2ClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)
+        with self.assertRaises(Exception) as context:
+            client.refresh_keys()
+        self.assertEqual(str(context.exception), 'Exception msg')
+
     @patch('uid2_client.client.parse_keys_json')
     def test_refresh_json(self, mock_refresh_keys, mock_parse_keys):
         mock_parse_keys.return_value = self._key_collection
@@ -33,7 +41,7 @@ def test_refresh_json(self, mock_refresh_keys, mock_parse_keys):
         self.assertIsNotNone(keys)
 
     def test_encrypt_decrypt(self, mock_refresh_keys_util):
-        mock_refresh_keys_util.return_value = self._key_collection
+        mock_refresh_keys_util.return_value = RefreshResponse.make_success(self._key_collection)
         client = Uid2ClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)
         client.refresh_keys()
 
@@ -45,7 +53,7 @@ def test_encrypt_decrypt(self, mock_refresh_keys_util):
         self.assertEqual(example_uid, result.uid)
 
     def test_can_decrypt_another_clients_encrypted_token(self, mock_refresh_keys_util):
-        mock_refresh_keys_util.return_value = self._key_collection
+        mock_refresh_keys_util.return_value = RefreshResponse.make_success(self._key_collection)
         client = Uid2ClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)
         client.refresh_keys()
 
@@ -58,7 +66,7 @@ def test_can_decrypt_another_clients_encrypted_token(self, mock_refresh_keys_uti
         self.assertEqual(example_uid, result.uid)
 
     def test_sharing_token_is_v4(self, mock_refresh_keys_util):
-        mock_refresh_keys_util.return_value = self._key_collection
+        mock_refresh_keys_util.return_value = RefreshResponse.make_success(self._key_collection)
         client = Uid2ClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)
         client.refresh_keys()
 
@@ -67,23 +75,23 @@ def test_sharing_token_is_v4(self, mock_refresh_keys_util):
         self.assertFalse(contains_base_64_special_chars)
 
     def test_uid2_client_produces_uid2_token(self, mock_refresh_keys_util):
-        mock_refresh_keys_util.return_value = self._key_collection
+        mock_refresh_keys_util.return_value = RefreshResponse.make_success(self._key_collection)
         client = Uid2ClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)
         client.refresh_keys()
 
         ad_token = client.encrypt(example_uid)
         self.assertEqual("A", ad_token[0])
 
     def test_euid_client_produces_euid_token(self, mock_refresh_keys_util):
-        mock_refresh_keys_util.return_value = self._key_collection
+        mock_refresh_keys_util.return_value = RefreshResponse.make_success(self._key_collection)
         client = EuidClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)
         client.refresh_keys()
 
         ad_token = client.encrypt(example_uid)
         self.assertEqual("E", ad_token[0])
 
     def test_raw_uid_produces_correct_identity_type_in_token(self, mock_refresh_keys_util):
-        mock_refresh_keys_util.return_value = self._key_collection
+        mock_refresh_keys_util.return_value = RefreshResponse.make_success(self._key_collection)
         client = Uid2ClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)
         client.refresh_keys()
 
@@ -102,7 +110,8 @@ def test_raw_uid_produces_correct_identity_type_in_token(self, mock_refresh_keys
 
     def test_multiple_keys_per_keyset(self, mock_refresh_keys_util):
         def get_post_refresh_keys_response_with_multiple_keys():
-            return create_default_key_collection([master_key, site_key, master_key2, site_key2])
+            return RefreshResponse.make_success(
+                create_default_key_collection([master_key, site_key, master_key2, site_key2]))
 
         mock_refresh_keys_util.return_value = get_post_refresh_keys_response_with_multiple_keys()
         client = Uid2ClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)
@@ -115,7 +124,7 @@ def get_post_refresh_keys_response_with_multiple_keys():
 
     def test_cannot_encrypt_if_no_key_from_default_keyset(self, mock_refresh_keys_util):
         def get_post_refresh_keys_response_with_no_default_keyset_key():
-            return create_default_key_collection([master_key])
+            return RefreshResponse.make_success(create_default_key_collection([master_key]))
 
         mock_refresh_keys_util.return_value = get_post_refresh_keys_response_with_no_default_keyset_key()
         client = Uid2ClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)
@@ -128,8 +137,8 @@ def get_post_refresh_keys_response_with_no_default_keyset_key():
     def test_cannot_encrypt_if_theres_no_default_keyset_header(self, mock_refresh_keys_util):
         def get_post_refresh_keys_response_with_no_default_keyset_header():
             key_set = [master_key, site_key]
-            return EncryptionKeysCollection(key_set, IdentityScope.UID2, 9000, 1,
-                                            "", 86400)
+            return RefreshResponse.make_success(EncryptionKeysCollection(key_set, IdentityScope.UID2, 9000, 1,
+                                            "", 86400))
 
         mock_refresh_keys_util.return_value = get_post_refresh_keys_response_with_no_default_keyset_header()
         client = Uid2ClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)
@@ -141,7 +150,7 @@ def get_post_refresh_keys_response_with_no_default_keyset_header():
 
     def test_expiry_in_token_matches_expiry_in_response(self, mock_refresh_keys_util):
         def get_post_refresh_keys_response_with_token_expiry():
-            return create_default_key_collection([master_key, site_key])
+            return RefreshResponse.make_success(create_default_key_collection([master_key, site_key]))
 
         mock_refresh_keys_util.return_value = get_post_refresh_keys_response_with_token_expiry()
         client = Uid2ClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)
@@ -168,7 +177,7 @@ def test_encrypt_key_inactive(self, mock_refresh_keys_util):
         def get_post_refresh_keys_response_with_key_inactive():
             inactive_key = EncryptionKey(245, site_id, now, now + dt.timedelta(days=1), now + dt.timedelta(days=2),
                                          site_secret, keyset_id=99999)
-            return create_default_key_collection([inactive_key])
+            return RefreshResponse.make_success(create_default_key_collection([inactive_key]))
 
         mock_refresh_keys_util.return_value = get_post_refresh_keys_response_with_key_inactive()
         client = Uid2ClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)
@@ -182,7 +191,7 @@ def test_encrypt_key_expired(self, mock_refresh_keys_util):
         def get_post_refresh_keys_response_with_key_expired():
             expired_key = EncryptionKey(245, site_id, now, now, now - dt.timedelta(days=1), site_secret,
                                         keyset_id=99999)
-            return create_default_key_collection([expired_key])
+            return RefreshResponse.make_success(create_default_key_collection([expired_key]))
 
         mock_refresh_keys_util.return_value = get_post_refresh_keys_response_with_key_expired()
         client = Uid2ClientFactory.create(self._CONST_BASE_URL, self._CONST_API_KEY, client_secret)