Code Testing with Flask, unittest and PyQuery

Over the last year I’ve been teaching programming.  The language of choice where i’ve been working is Python (which I had not used very much before going in).  Because of that, I’ve been making it a point to code as many of my personal projects as I can in the language.  Despite a long “love/hate” relationship over the first few months (python was my first language that uses syntactical whitespace), I’ve really grown to love it.

I’ve also become quite fond of the Flask framework for developing web applications.  I love its minimal approach to things, really getting out of the way of things that a web framework doesn’t need much business sticking its nose in.

One project that I’ve been working on a lot lately has started to grow to the point where I really started feeling guilty about not having written any tests.  It’s that classic, “well I intend to write test real soon now…”.  We cover a small amount of the unittest module where I teach, but I felt I really needed to spend some time diving into how to write some good tests for a Flask-based application.

I spent a fair amount of time over the last week searching the internet for good techniques to do unit and functional testing using Flask with varying degrees of success.  There’s lots of material out there, but no one source covered everything that I wanted.

  • I wanted to write unit and functional/integration tests.  So much of the code in my application is in Flask views, Jinja2 Templates and SQLAlchemy models, unit testing doesn’t make sense for me.
  • To properly test my code, I want to run tests against either a test database or mock.  I haven’t yet found a good mockup, so I’ve been focused running against a test database for now.  The script to run the test would need to handle managing the schema and test data.
  • I want to make room for me to write a lot of tests, so I wanted my tests to be able to cleanly divide into multiple files.  My application has some clear division lines for functionally and I want to divide my test suites against those lines

So here’s what i’ve come up with so far.  This is me taking the good parts from what I’ve liked over the web (things like the wonderful Flask Mega Tutorial have been an invaluable starting point).  I have not found examples of anyone using PyQuery in the way I have been, so maybe I’m bringing a new twist on some things.

Structure

Before we dive in, a quick overview of how I’ve structured my web app so far.  My directory structure looks something like below.  This isn’t a complete list of all files, but at least some of the interesting/relevant ones.

  • MyProject/     (app root folder)
    • alembic/    (schema versioning)
    • app/     (main application code)
      • forms/   (WTForms files, divided by site functional areas)
      • helpers.py   (misc/stand alone utility functions with no better place to go)
      • model.py   (SQLAlchemy classes for data models)
      • static/    (static files)
      • templates/  (jinja2 templates for views)
      • views.py  (Flask routes)
    • config.py   (configuration settings.)
    • env/    (virtualenv environment)
    • run_tests.py    (Main script to call test suite)
    • tests/   (test suite files)
    • web.py   (main script used to run the Flask application)

So many examples on the web only talk about writing tests in the context of one file.  It was important to me to be able to organize all my test scripts in a single directory and be able to easily run them from one place.

Some sites showed examples where the test files were in a directory, but to launch them you had to invoke python and remember to include a module on the command line with some crazy command line options or some nonsense.  Ain’t nobody got time to remember all that.  I wanted one script to run everything, or the option to run any individual script directly.

Tools

I went with unittest for the main framework.  It’s simple and non-complicated.  I saw lots of sites speak highly of nose, but unittest has not caused me any trouble so far to warrant spending the time trying something different.  The TestSuite class also let me group by tests just like I wanted to.

I didn’t go looking for coverage, it found me.  It was so easy to add it was a no-brainer.  Now I get a nice report at the end of my tests letting me know how much code my tests are actually covering.  I haven’t looked at the html report yet but it makes me happy just to know its there.

For testing the views, most sites I found were just doing a simple string search on the resulting html from the rendered template.  Something like:

assert "First Name" in page

Yeah, this works…but it’s way too not specific enough for me.  If I post to a form, I want to check did a specific alert <div> appear with the right message, or that the form field was given an error class so it would be highlighted in red.  For that, it meant I would turn to PyQuery.  BeautifulSoup would also do the same job, but I really like the ease of using the jquery-style selectors.

Putting It All Together

Ok, enough with the backstory.  Show me some freakin’ code. I’ve posted this code on a gist as well if you prefer to read things there.

First up, the main script run_tests.py. It doesn’t have any test cases itself, but calls everything to make things run. It uses the TestSuite runner instead of unittest.main() to run.

RUN_TESTS.PY

import unittest

# Import Test Suite Classes
from tests.mainsite_tests import MainSiteTest
from tests.manager_site_tests import ManagerSiteTest
from tests.myapp_unittests import MyAppUnitTest

# Track Code Coverage Report
from coverage import coverage
cov = coverage(branch = True, omit = ['env/*', 'run_tests.py', 'tests/*'])
cov.start()

# Returns a TestSuite for all the test classes we want to run
def suite():
    suite = unittest.TestSuite()

    suite.addTest(unittest.makeSuite(MyAppUnitTest))
    suite.addTest(unittest.makeSuite(MainSiteTest))
    suite.addTest(unittest.makeSuite(ManagerSiteTest))

    return suite

# If we run from the command line directly, run the test suite
if __name__ == '__main__':
    runner = unittest.TextTestRunner()
    test_suite = suite()

    # Catch exceptions in our tests so our teardown functions can run
    try:
        runner.run(test_suite)
    except:
        pass
    
    # Display the Code Coverage Report
    cov.stop()

    # I'm not currently saving the report, but I can change my mind in the future
    #cov.save()
    
    print "\n\nCoverage Report:\n"
    cov.report()
    
    # I'm also not currently using the HTML version, but may want to later
    #print "HTML version: " + os.path.join(basedir, 'tmp/coverage/index.html')
    #cov.html_report(directory='tmp/coverage')
    cov.erase()

First test script is the one for my unit tests. Most of my code is in the MVC classes, but I do have some independent functions and general utility type functions that I keep in a file called helpers.py. Most of these are deterministic functions, perfect for unit testing, so I started my test writing there. Things like date manipulation functions (I like all my dates to be in epoch time). The content of the actual tests are really the focus of this post, but I’m including a couple tests in here for completeness.

MYAPP_UNITTESTS.PY

# If we want to call this script directly, we need to add the parent
# directory to our search path so we can find everything
if __name__ == '__main__':
    import os, sys
    basedir = os.path.abspath(os.path.dirname(__file__) + '/../')
    sys.path.insert(0, basedir)

import unittest
import datetime

# app is my Flask class, db is the sqlalchemy db connection
from app import app, db
from app.helpers import *

class MyAppUnitTest(unittest.TestCase):

    def test_example(self):
    self.failUnless( 1+1 == 2, 'one plus one fails!')
    
    # date_time_to_epoch converts a Python datetime variable to an epoch timestamp (int)
    def test_datetime_to_epoch(self):
    dt = datetime.datetime(2013, 8, 8, 8, 0, 0)
    self.assertEqual( datetime_to_epoch(dt, 'America/New_York'), 1375963200.0)
    
    # epoch_to_datetime converts an epoch timestamp (int) to a Python datetime
    def test_epoch_to_datetime(self):
    dt = epoch_to_datetime(1375963200, 'America/New_York')
    
    self.assertEqual(dt.year, 2013)
    self.assertEqual(dt.month, 8)
    self.assertEqual(dt.day, 8)
    self.assertEqual(dt.hour, 8)
    self.assertEqual(dt.minute, 0)
    
    # select_choices_dict takes a list of tuples (used to generate select menus)
    # and returns a dictionary of key/value pairs for easy lookups.
    def test_select_choices_dict(self):
    menu = [
        ('pending', 'Pending'),
        ('open', 'Open'),
        ('closed', 'Closed'),
        ('completed', 'Completed')
    ]
    
    d = select_choices_dict(menu)
    
    self.assertEqual(d['pending'], 'Pending')
    self.assertEqual(d['open'], 'Open')
    self.assertEqual(d['closed'], 'Closed')
    self.assertEqual(d['completed'], 'Completed')
    
    # If we call this script directly on the command line, run the unittests.
    # Normally we'll be called from run_tests.py
    if __name__ == '__main__':
        try:
            unittest.main()
        except:
            pass

Now with the unit tests out of the way, we can get to the fun stuff — the functional tests. For this, we’ll call our Flask app just like a browser would and test the resulting page content and status code. Thankfully, Flask comes with a Test Client built in that lets us GET and POST to our little hearts content. All we have to do it tell Flask we’re in “testing” mode.

So now all we have left to deal with is how we’re going to handle our database.

Its very important when performing these kinds of tests you not only run them against a dedicated test database but that also each test run must use a known good starting state of the database. This means each time you run the tests they must clean up after themselves. It’s no good if you’ve got a test to insert a user into the system if you’ve got a unique requirement in your database.

I found lots of pointers to framework techniques to to query/transaction isolation on your test cases — basically do a “rollback” on your whole test suite at the end like it never ran in the first place. I found the description of these to be overly complicated. For me, just building a clean database from scratch wouldn’t take that long and seems like the least convoluted solution. This means we’ll just call SQLAlchemy’s create_all() to create the schema, and drop_all() to clean up after ourselves.

Pretty much every example I found had you doing create_all() at setup time and drop_all() at teardown. I want to be able to go into the database and look at what the test suite did, so I didn’t like having drop_all() in the teardown, so I just run it before the create_all() in setup.

There are two types of setup/teardown commands you can use with unittest. The setUp() and tearDown() methods run before and after each individual test function. That is handy, but I don’t want to put my database setup methods in there since it will cause the entire database to be recreated for each test case. You may hear some people say that’s the only proper way to do testing, but it’s not what I’m looking for out of my tests. I want to be able to have my tests simulate a user interacting with my application from setup through usage of each feature. Rather than trying to build up a database full of mock data to cleanly use for each test, I want to use the data created from the successful run of one test flow into the next.

The second pair of setup/teardown commands are setUpClass() and tearDownClass(). They run before and after all the tests in a given class run and that’s what I’ve elected to place my database setup code. As I mentioned early, I’m dividing my test classes up by functional sections of the application, so having a clean database for each of those sections is perfect.

Yes, this is not “clean” or “pure”. Yes, a failure in an early test could cause the rest of the tests in the suite to fail. Since my project isn’t *that* big and I’m the primary developer working on it, if I break a test it’s me who has to deal with it, so that’s a conscious decision I have made. It may not work in in your scenario. Hell, I could change my mind down the road and move to mockups and stubs. But for now, I’m liking this solution.

To make this all happen, I first create a base class that inherits from unittest.TestCase. Each class that inherits from this class will share the setUpClass() and tearDownClass() methods that will configure Flask in testing mode and setup the database.

MYAPP_TEST.PY

import os
basedir = os.path.abspath(os.path.dirname(__file__))

import unittest
# app contains Flask, db is SQLAlchemy
from app import app, db

# Parent class for our functional tests
class MyAppTest(unittest.TestCase):
    @classmethod
    def setUpClass(self):
        # Put Flask into TESTING mode for the TestClient
        app.config['TESTING'] = True
        # Disable CSRF checking for WTForms
        app.config['WTF_CSRF_ENABLED'] = False
        # Point SQLAlchemy to a test database location
        # (set in virtualenv normally, but fall back to sqlite if not defined)
        app.config['SQLALCHEMY_DATABASE_URI'] = os.environ.get('TEST_DATABASE_URI', 'sqlite:///' + os.path.join(basedir, 'hakaru_test.db'))
        
        # Drop and Re-Create the Schema
        db.drop_all()
        db.create_all()
    
        # If you need to initialize any data in your tables before your app can run
        # do that here. I don't need to for my app.
    
    @classmethod
    def tearDownClass(self):
        db.session.remove()
        # I moved the drop to setup, but you usually see it here
        # db.drop_all()
    
    def setUp(self):
        # Get a new instance of the Flask Test Client for each test
        self.app = app.test_client()

And finally, the file with the actual functional tests. I have multiple of these, but I’m just showing one here to demonstrate the point.

MAINSITE_TESTS.PY

# If we want to call this script directly, we need to add the parent
# directory to our search path so we can find everything
if __name__ == '__main__':
    import os, sys
    basedir = os.path.abspath(os.path.dirname(__file__) + '/../')
    sys.path.insert(0, basedir)

from app import app, db
from myapp_test import MyAppTest
from pyquery import PyQuery as pq

from app.model import *

class MainSiteTest(MyAppTest):
    
    # Test that the main site index page loads.
    # In this test, I just care that we got a 200 OK status code
    # returned, I don't care as much about the content of the page.
    def test_index(self):
        rv = self.app.get('/')
        
        # Make sure we got an OK response
        self.assertEqual(rv.status_code, 200)
    
    # A thorough test of the signup page.
    def test_signup(self):
    # First test viewing the page returns a 200 OK status code
    rv = self.app.get('/signup')
    
    # Make sure we got an OK response
    self.assertEqual(rv.status_code, 200)
    
    # Next, if we try to POST a form with empty data, we should
    # trigger a whole slew of validation errors.
    data = {}
    
    # Posting bad data should still generate a 200 OK
    rv = self.app.post('/signup', data=data)
    self.assertEqual(rv.status_code, 200)
    
    # Validation errors will be "flash" messages to the user
    # that will show up in div's with the "alert" class.
    # Use PyQuery to find those elements and check that
    # the expected field names are mentioned in the error
    # messages.
    q = pq(rv.data)
    alerttext = q('.alert').text()
    
    assert 'First Name' in alerttext
    assert 'Last Name' in alerttext
    assert 'Organization Name' in alerttext
    assert 'Invitation Code' in alerttext
    assert 'E-Mail Address' in alerttext
    assert 'Password' in alerttext
    assert 'Accept TOS' in alerttext
    
    # Now submitting with "real" data should not produce errors
    # and should result in new database entries in three tables
    data = {
        'invite_code': 'AskMeAboutGofers',
        'email': 'testuser@test.com',
        'orgname': 'My Test Organization',
        'givenname': 'Test',
        'surname': 'User',
        'password': 'qwerty123',
        'confirm': 'qwerty123',
        'accept_tos': 'y'
    }
    
    rv = self.app.post('/signup', data=data)
    self.assertEqual(rv.status_code, 302) # On success we redirect
    
    # Check that database entries were created
    # First up is a user record
    u = User.query.filter_by(email=data['email']).all()
    
    self.assertEqual(len(u), 1, 'Incorrect number of users created')
    # Test that the data was saved in the database correctly
    self.assertEqual(u.givenname, data['givenname'])
    self.assertEqual(u.surname, data['surname'])
    self.assertEqual(u.email, data['email'])
    
    # Signup also creates an organization record
    o = Organization.query.filter_by(name='My Test Organization').all()
    
    self.assertEqual(len(o), 1, 'Incorrect number of organizations created')
    self.assertEqual(o.name, data['orgname'])
    
    # OrganizationUsers connects the user to the organization
    ou = OrganizationUsers.query.filter_by(organization_id=o[0].id, user_id=u[0].id, role='owner').all()
    
    self.assertEqual(len(ou), 1, 'Incorrect number of organization user records created')

if __name__ == '__main__':
    try:
        unittest.main()
    except:
        pass

Conclusion

I really like the completeness of these tests. It seems silly to me to write a bunch of code just to test one component like a model. I have a fairly high degree of confidence that SQLAlchemy will save my data to the database, I don’t need to test that. I do need to test that if I send data view my web form that it will actually make it where its supposed to go.

Sure, this doesn’t get everything. There’s no tests in here for things that are really view specific (like Javascript). But it does test some of the most important parts of my web app, the code in Flask.

You Might Also Like