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()