Fixed dupicite key error

auth_token/exceptions.py

@@ -0,0 +1,2 @@
+class KeyGeneratorError(Exception):
+    pass

auth_token/models.py

@@ -13,6 +13,7 @@
 
 from enumfields import NumEnumField
 
+from chamber.exceptions import PersistenceException
 from chamber.models import SmartModel, SmartQuerySet, SmartManager
 
 from generic_m2m_field.models import GenericManyToManyField
@@ -21,6 +22,7 @@
 from auth_token.config import settings
 
 from .enums import AuthorizationRequestState, AuthorizationRequestResult
+from .exceptions import KeyGeneratorError
 
 
 KEY_SALT = 'django-auth-token'
@@ -40,16 +42,22 @@ def _hash_key(self, key, **kwargs):
         return hash_key(key)
 
     def create(self, key_generator, **kwargs):
-        for attempt in range(settings.MAX_RANDOM_KEY_ITERATIONS + 1):
+        for attempt in range(1, settings.MAX_RANDOM_KEY_ITERATIONS + 1):
             try:
                 key = key_generator()
                 hashed_key = self._hash_key(key, **kwargs)
                 obj = super().create(key=hashed_key, **kwargs)
                 obj.secret_key = key
                 return obj
             except IntegrityError:
-                if attempt > settings.MAX_RANDOM_KEY_ITERATIONS:
-                    raise IntegrityError('Could not produce unique key')
+                if attempt >= settings.MAX_RANDOM_KEY_ITERATIONS:
+                    raise KeyGeneratorError('Could not produce unique key')
+            except PersistenceException as ex:
+                if ex.error_dict and 'key' in ex.error_dict:
+                    if attempt >= settings.MAX_RANDOM_KEY_ITERATIONS:
+                        raise KeyGeneratorError('Could not produce unique key')
+                else:
+                    raise ex
 
 
 class AuthorizationTokenManager(BaseHashKeyManager):

example/Makefile

@@ -40,9 +40,7 @@ pip:
 	$(PYTHON_BIN)/pip install -r requirements.txt
 
 initvirtualenv:
-	virtualenv -p $(PYTHON) --no-site-packages $(VIRTUAL_ENV)
-	$(PYTHON_BIN)/pip install --upgrade pip==1.5
-	$(PYTHON_BIN)/pip install setuptools --no-use-wheel --upgrade
+	virtualenv -p $(PYTHON) $(VIRTUAL_ENV)
 
 bootstrap: initvirtualenv pip
 

example/dj/apps/app/tests/utils.py

@@ -19,6 +19,7 @@
 )
 
 from auth_token.enums import AuthorizationRequestState
+from auth_token.exceptions import KeyGeneratorError
 from auth_token.models import AnonymousUser, AnonymousAuthorizationToken, MobileDevice
 from auth_token.utils import (
     compute_expires_at, hash_key, header_name_to_django, generate_key,
@@ -698,3 +699,23 @@ def test_cancel_authorization_request_should_cancel_authorization_and_not_call_r
             cancel_authorization_request(authorization_request)
             mocked_receiver.assert_not_called()
             assert_equal(authorization_request.refresh_from_db().state, AuthorizationRequestState.CANCELLED)
+
+    @freeze_time(now())
+    def test_duplicite_key_generator_should_raise_exception(self):
+        create_otp('test', key_generator=lambda: '1234')
+        with assert_raises(KeyGeneratorError):
+            create_otp('test', key_generator=lambda: '1234')
+
+    @freeze_time(now())
+    def test_key_genrator_should_be_called_more_times(self):
+        def key_genertor():
+            key_genertor.iteration += 1
+            if key_genertor.iteration < 100:
+                return '1234'
+            else:
+                return '1235'
+        key_genertor.iteration = 0
+
+        create_otp('test', key_generator=lambda: '1234')
+        assert_equal(create_otp('test', key_generator=key_genertor).secret_key, '1235')
+        assert_equal(key_genertor.iteration, 100)

example/requirements.txt

@@ -2,7 +2,7 @@ Django==3.1
 django-germanium==2.2.5
 coverage==5.3.1
 django-pyston==2.14.2
-django-is-core==2.22.3
+django-is-core==2.22.9
 djangorestframework==3.12.2
 freezegun==1.0.0
 -e ../

setup.py

@@ -27,7 +27,7 @@
         'django>=2.2.14',
         'django-ipware>=3.0.2',
         'import_string==0.1.0',
-        'django-chamber>=0.5.26',
+        'django-chamber>=0.6.5',
         'django-generic-m2m-field>=0.0.4',
         'django-choice-enumfields>=1.0.3',
     ],