Chapter 15: Unit tests with the Flask test client (15b)

miguelgrinberg · miguelgrinberg · commit d305c3d28b14 · 2015-11-05T11:09:17.000-08:00
diff --git a/config.py b/config.py
@@ -32,6 +32,7 @@ class TestingConfig(Config):
     TESTING = True
     SQLALCHEMY_DATABASE_URI = os.environ.get('TEST_DATABASE_URL') or \
         'sqlite:///' + os.path.join(basedir, 'data-test.sqlite')
+    WTF_CSRF_ENABLED = False
 
 
 class ProductionConfig(Config):
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -0,0 +1,54 @@
+import re
+import unittest
+from flask import url_for
+from app import create_app, db
+from app.models import User, Role
+
+class FlaskClientTestCase(unittest.TestCase):
+    def setUp(self):
+        self.app = create_app('testing')
+        self.app_context = self.app.app_context()
+        self.app_context.push()
+        db.create_all()
+        Role.insert_roles()
+        self.client = self.app.test_client(use_cookies=True)
+
+    def tearDown(self):
+        db.session.remove()
+        db.drop_all()
+        self.app_context.pop()
+
+    def test_home_page(self):
+        response = self.client.get(url_for('main.index'))
+        self.assertTrue(b'Stranger' in response.data)
+
+    def test_register_and_login(self):
+        # register a new account
+        response = self.client.post(url_for('auth.register'), data={
+            'email': 'john@example.com',
+            'username': 'john',
+            'password': 'cat',
+            'password2': 'cat'
+        })
+        self.assertTrue(response.status_code == 302)
+
+        # login with the new account
+        response = self.client.post(url_for('auth.login'), data={
+            'email': 'john@example.com',
+            'password': 'cat'
+        }, follow_redirects=True)
+        self.assertTrue(re.search(b'Hello,\s+john!', response.data))
+        self.assertTrue(
+            b'You have not confirmed your account yet' in response.data)
+
+        # send a confirmation token
+        user = User.query.filter_by(email='john@example.com').first()
+        token = user.generate_confirmation_token()
+        response = self.client.get(url_for('auth.confirm', token=token),
+                                   follow_redirects=True)
+        self.assertTrue(
+            b'You have confirmed your account' in response.data)
+
+        # log out
+        response = self.client.get(url_for('auth.logout'), follow_redirects=True)
+        self.assertTrue(b'You have been logged out' in response.data)
