1 import unittest

     2 

     3 from iro.model import schema

     4 from iro.model.utils import WithSession

     5 

     6 class DBTestCase(unittest.TestCase):

     7     '''a TestCase with DB connection

     8     you have to set self.session manually in setUp'''

     9     def __init__(self,*args,**kwargs):

    10         self.engine=None

    11         unittest.TestCase.__init__(self,*args,**kwargs)

    12 

    13     def tearDown(self):

    14         self.__cleanDB()

    15     

    16     def session(self,autocommit=True):

    17         '''returns a WithSession'''

    18         return WithSession(self.engine,autocommit)

    19     

    20     def __cleanDB(self):

    21         '''cleaning database'''

    22         with self.session() as session:

    23             for table in schema.__tables__:

    24                 session.query(getattr(schema,table)).delete()