Adds validation to force creators to be a member of the Team. This is a partial fix for [#2].

goatctf/core/models.py

@@ -1,5 +1,6 @@
 from django.contrib.auth.models import User
 from django.db import models
+from django.db.utils import IntegrityError
 import markdown
 
 from core.settings import CHALLENGE_NAME_LENGTH, FLAG_LENGTH, TEAM_NAME_LENGTH
@@ -40,6 +41,15 @@ class Team(models.Model):
     name = models.CharField(max_length=TEAM_NAME_LENGTH, unique=True)
     creator = models.ForeignKey("Player", related_name="created_teams")
 
+    def save(self, *args, **kwargs):
+        if not hasattr(self, 'creator'):
+            raise IntegrityError("Creator must be defined.")
+        if self.creator.team and self.creator.team != self:
+            raise IntegrityError("Creator must be a part of the team")
+        super(Team, self).save(*args, **kwargs)
+        self.creator.team = self
+        self.creator.save()
+
     def __str__(self):
         return self.name
 
@@ -48,6 +58,11 @@ class Player(User):
     """A player is a user with a team."""
     team = models.ForeignKey("Team", blank=True, null=True)
 
+    def save(self, *args, **kwargs):
+        teams_exist = self.team and self.created_teams.count() > 0 
+        if teams_exist and self.team != self.created_teams.all()[0]:
+            raise IntegrityError("Player must be a part of all created teams!")
+        super(Player, self).save(*args, **kwargs)
 
 class Solution(models.Model):
     """A solution is a player's """

goatctf/tests/test_player.py

@@ -0,0 +1,34 @@
+from django.db.utils import IntegrityError
+import pytest
+
+from core.models import Player, Team
+
+
+@pytest.fixture
+def player_factory():
+    class PlayerFactory(object):
+        def get(self):
+            user = Player()
+            i=0
+            while Player.objects.filter(username=str(i)).count() > 0:
+                i+=1
+            user.username = str(i)
+            user.password = ''
+            user.save()
+            return user
+    return PlayerFactory()
+
+
+@pytest.mark.django_db
+def test_creator_team_equal(player_factory):
+    player1 = player_factory.get()
+    team1 = Team(name="Team 1", creator=player1)
+    team1.save()
+
+    player2 = player_factory.get()
+    team2 = Team(name="Team 2", creator=player2)
+    team2.save()
+
+    player1.team = team2
+    with pytest.raises(IntegrityError):
+        player1.save()

goatctf/tests/test_team.py

@@ -5,26 +5,48 @@ from core.models import Player, Team
 
 
 @pytest.fixture
-def user():
-    user = Player()
-    user.username = 'user'
-    user.password = ''
-    user.save()
-    return user
+def player_factory():
+    class PlayerFactory(object):
+        def get(self):
+            user = Player()
+            # FIXME There must be a cleaner way to do this.
+            i = 0
+            while Player.objects.filter(username=str(i)).count() > 0:
+                i += 1
+            user.username = str(i)
+            user.password = ''
+            user.save()
+            return user
+    return PlayerFactory()
 
 
 @pytest.mark.django_db
-def test_cannot_create_without_creator(user):
+def test_cannot_create_without_creator():
     team1 = Team(name="Team 1")
     with pytest.raises(IntegrityError):
         team1.save()
 
 
 @pytest.mark.django_db
-def test_team_names_are_unique(user):
+def test_team_names_are_unique(player_factory):
+    user1 = player_factory.get()
+    team1 = Team(name="Team 1", creator=user1)
+    team1.save()
+
+    user2 = player_factory.get()
+    team2 = Team(name="Team 1", creator=user2)
+    with pytest.raises(IntegrityError):
+        team2.save()
+
+
+@pytest.mark.django_db
+def test_creators_are_unique(player_factory):
+    user = player_factory.get()
+
     team1 = Team(name="Team 1", creator=user)
     team1.save()
 
-    team2 = Team(name="Team 1", creator=user)
+    team2 = Team(name="Team 2", creator=user)
     with pytest.raises(IntegrityError):
         team2.save()
+