Use refresh_from_db

dimagi · Sep 2, 2024 · c50212b · c50212b
1 parent 8b7b7cb
commit c50212b
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/corehq/apps/sso/tests/test_tasks.py b/corehq/apps/sso/tests/test_tasks.py
@@ -253,8 +253,8 @@ def test_updates_non_compliant_expiry_to_max(self):
         with freeze_time(current_time):
             num_updates = enforce_key_expiration_for_idp(idp)
 
-        updated_key = HQApiKey.objects.get(id=key.id)
-        self.assertEqual(updated_key.expiration_date, current_time + datetime.timedelta(days=30))
+        key.refresh_from_db()
+        self.assertEqual(key.expiration_date, current_time + datetime.timedelta(days=30))
         self.assertEqual(num_updates, 1)
 
     def test_exits_when_no_maximum_date_exists(self):
@@ -288,8 +288,8 @@ def test_integration_test(self):
         with freeze_time(current_time):
             update_sso_user_api_key_expiration_dates(idp.id)
 
-        updated_key = HQApiKey.objects.get(id=key.id)
-        self.assertEqual(updated_key.expiration_date, current_time + datetime.timedelta(days=30))
+        key.refresh_from_db()
+        self.assertEqual(key.expiration_date, current_time + datetime.timedelta(days=30))
 
     def _create_idp_for_domains(self, domains, max_days_until_user_api_key_expiration=30):
         idp = generator.create_idp('test-idp', account=self.account)