david / django-invitation-backend / changeset / be8229c1326b

@@ -50,9 +50,12 @@ class InvitationKeyManager(models.Manage
     def remaining_invitations_for_user(self, user):
         """
-        Returns the number of remaining invitations for a given ``User``.
+        Return the number of remaining invitations for a given ``User``.
         """
-        return InvitationUser.objects.get(inviter=user).invitations_remaining
+        invitation_user, created = InvitationUser.objects.get_or_create(
             inviter=user,
             defaults={'invitations_remaining': settings.INVITATIONS_PER_USER})
         return invitation_user.invitations_remaining
     def delete_expired_keys(self):
         for key in self.all():

@@ -121,6 +121,15 @@ class InvitationModelTests(InvitationTes
         remaining = remaining_invites(self.sample_user)
         self.assertEqual(remaining, new_remaining)
         # If no InvitationUser (for pre-existing/legacy User), one is created
         old_sample_user = User.objects.create_user(username='lewis',
                                                    password='secret',
                                                    email='lewis@example.com')
         old_sample_user.invitationuser_set.all().delete()
         self.assertEqual(old_sample_user.invitationuser_set.count(), 0)
         remaining = remaining_invites(old_sample_user)
         self.assertEqual(remaining, settings.INVITATIONS_PER_USER)
 class InvitationFormTests(InvitationTestCase):
     """