From 5c1d21f91eb05538707fb5a31c38f4a271000949 Mon Sep 17 00:00:00 2001
From: Di Wu <dw323@cornell.edu>
Date: Sun, 15 Nov 2015 15:58:48 -0800
Subject: [PATCH] make welcome code logic resilient

---
 hc/front/tests/test_basics.py | 13 +++++++++++++
 hc/front/views.py             | 10 ++++++----
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/hc/front/tests/test_basics.py b/hc/front/tests/test_basics.py
index 72b0991c..86425a79 100644
--- a/hc/front/tests/test_basics.py
+++ b/hc/front/tests/test_basics.py
@@ -1,4 +1,5 @@
 from django.test import TestCase
+from hc.api.models import Check
 
 
 class BasicsTestCase(TestCase):
@@ -6,3 +7,15 @@ class BasicsTestCase(TestCase):
     def test_it_shows_welcome(self):
         r = self.client.get("/")
         self.assertContains(r, "Get Notified", status_code=200)
+
+    def test_welcome_code(self):
+        r = self.client.get("/")
+        code = self.client.session["welcome_code"]
+        assert Check.objects.filter(code=code).exists()
+
+        self.client.session["welcome_code"] = "x"
+        r = self.client.get("/")
+        code = self.client.session["welcome_code"]
+        assert r.status_code == 200
+        assert code != "x"
+        assert Check.objects.filter(code=code).exists()
diff --git a/hc/front/views.py b/hc/front/views.py
index 4f1de064..afe89586 100644
--- a/hc/front/views.py
+++ b/hc/front/views.py
@@ -27,14 +27,16 @@ def index(request):
     if request.user.is_authenticated():
         return redirect("hc-checks")
 
-    if "welcome_code" not in request.session:
+    check = None
+    if "welcome_code" in request.session:
+        code = request.session["welcome_code"]
+        check = Check.objects.filter(code=code).first()
+
+    if check is None:
         check = Check()
         check.save()
         code = str(check.code)
         request.session["welcome_code"] = code
-    else:
-        code = request.session["welcome_code"]
-        check = Check.objects.get(code=code)
 
     ctx = {
         "page": "welcome",