import unittest

from iro.model import schema

from iro.model.utils import WithSession

class DBTestCase(unittest.TestCase):

    '''a TestCase with DB connection

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

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

        self.engine=None

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

    def tearDown(self):

        self.__cleanDB()

    def session(self,autocommit=True):

        '''returns a WithSession'''

        return WithSession(self.engine,autocommit)

    def __cleanDB(self):

        '''cleaning database'''

        with self.session() as session:

            for table in schema.__tables__:

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