Flask - (7)

Database, ORM, creating models, testing python

2020-01-08
by cS


For this occasion we will use a third party libary for data models, peewee, to install peewee first download this dependency from pip.

pip install peewee

I will name my storage file, the file containing models and db information storage.py, additionally I will place this storage.py file inside the folder where the app.py file is located, the root application folder.

from peewee import *

TESTING = True # Gives logic.

# The config dictionary, for meta.
STORE_CONFIG = { 
    "debug_db" : SqliteDatabase(
                    "units.db",
                    pragmas={"journal_mode": "wal"}
                ),
    "test_db" : SqliteDatabase(
                    "test_units.db",
                    pragmas={'journal_mode': 'wal'}
                )
}


class BaseModel(Model):
    class Meta:
        if TESTING:
            database = STORE_CONFIG["test_db"]
        else:
            database = STORE_CONFIG["debug_db"]


class Unit(BaseModel):
    area_of_operation = CharField(max_length=100, default="Unknown")
    unit_name         = CharField(max_length=100, unique=True)
    location          = CharField(max_length=100, default="Unknown")
    unit_count        = IntegerField(default=13)
    casualties        = IntegerField(default=0)


def initialize(dbs):
    if dbs.is_closed() == False:
        dbs.close()
    dbs.connect()
    dbs.create_tables([Unit], safe=True)
    dbs.close()

We should test this. I made a file called test.py file also at the root directory.

import os 
import unittest

from peewee import *

from storage import initialize, Unit, STORE_CONFIG
    
THIS_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)))

class BaseTest(unittest.TestCase):
    def setUp(self):
        self.db = "test_units.db"
        self.tu = {
            "unit_name" : "CWSTF6",
            "unit_count": 12,
            "casualties": 1,
        }
    
    def tearDown(self): # delete the db file on each test function.
        if os.path.exists(os.path.join(THIS_PATH, self.db)):
            os.remove(os.path.join(THIS_PATH, self.db))


# Inherited the fixture from from Base Test
class TestUnitModel(BaseTest):
    def test_create_unit(self): # creates a unit, and assets it can be retreived.
        initialize(STORE_CONFIG["test_db"])
        unit = Unit(**self.tu)
        unit.save()
        self.assertEqual(Unit.get_by_id(unit.id), unit)
        with self.assertRaises(IntegrityError): # asserts integrity error from constraint.
            Unit(**self.tu).save()


if __name__ == "__main__":
    unittest.main()