Fix deleting a team with flags (#84)

mgerst · Oct 28, 2019 · 6f871f1 · 6f871f1
1 parent 236e0a3
commit 6f871f1
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/flag_slurper/autolib/models.py b/flag_slurper/autolib/models.py
@@ -67,7 +67,7 @@ def __str__(self):
 
 class Flag(BaseModel):
     id = peewee.AutoField(primary_key=True)
-    team = peewee.ForeignKeyField(Team, backref='flags')
+    team = peewee.ForeignKeyField(Team, backref='flags', on_delete='CASCADE')
     name = peewee.CharField(max_length=150)
 
 

diff --git a/tests/autolib/test_models.py b/tests/autolib/test_models.py
@@ -1,7 +1,6 @@
 import click
-import pytest
 
-from flag_slurper.autolib.models import CredentialBag, Credential, CaptureNote
+from flag_slurper.autolib.models import CaptureNote, Flag
 
 SUDO_FLAG = click.style('!', fg='red', bold=True)
 
@@ -38,3 +37,10 @@ def test_capture_note__str__(flag, service):
 def test_capture_note_sudo__str__(flag, service):
     note = CaptureNote(flag=flag, service=service, data='abcd', location='/root/test.flag', notes='did stuff\nUsed Sudo')
     assert note.__str__() == "/root/test.flag -> abcd{}".format(SUDO_FLAG)
+
+
+def test_flag_on_delete_cascade(team, flag):
+    initial = Flag.select().count()
+    team.delete().execute()
+    after = Flag.select().count()
+    assert initial - after == 1