From 9b63a06bfcd51b86ab20f91c48088046bfb5d68f Mon Sep 17 00:00:00 2001 From: Ned Batchelder Date: Thu, 26 Apr 2018 15:50:05 -0400 Subject: [PATCH] Properly clean up a wrapped function --- .../commands/tests/test_clear_expired_tokens.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/openedx/core/djangoapps/oauth_dispatch/management/commands/tests/test_clear_expired_tokens.py b/openedx/core/djangoapps/oauth_dispatch/management/commands/tests/test_clear_expired_tokens.py index 8ba8bc8776..cdd907389d 100644 --- a/openedx/core/djangoapps/oauth_dispatch/management/commands/tests/test_clear_expired_tokens.py +++ b/openedx/core/djangoapps/oauth_dispatch/management/commands/tests/test_clear_expired_tokens.py @@ -25,7 +25,7 @@ def counter(fn): """ def _counted(*largs, **kargs): _counted.invocations += 1 - fn(*largs, **kargs) + return fn(*largs, **kargs) _counted.invocations = 0 return _counted @@ -90,8 +90,11 @@ class EdxClearExpiredTokensTests(TestCase): AccessToken.objects.filter(refresh_token__isnull=True, expires__lt=now).count(), initial_count ) + original_delete = QuerySet.delete QuerySet.delete = counter(QuerySet.delete) - - call_command('edx_clear_expired_tokens', batch_size=1, sleep_time=0) - self.assertEqual(QuerySet.delete.invocations, initial_count) - self.assertEqual(AccessToken.objects.filter(refresh_token__isnull=True, expires__lt=now).count(), 0) + try: + call_command('edx_clear_expired_tokens', batch_size=1, sleep_time=0) + self.assertEqual(QuerySet.delete.invocations, initial_count) + self.assertEqual(AccessToken.objects.filter(refresh_token__isnull=True, expires__lt=now).count(), 0) + finally: + QuerySet.delete = original_delete