CODEWITHSUNDEEP
pythonpython-unittestpython-decorators
Alex
Alex

Reputation: 6047

how to pass arguments to a decorator in a class

I have a unit test that loads some data in the database to test against. I have created a decorator function for this, to be used on a test case

class TestCompany(unittest.TestCase):
    def setUp(self):
        ...
        self.engine = create_engine(db_url)
        Session = sessionmaker(bind=self.engine)
        self.session = Session() 


    @loadfixtures(['employee.json'])
    def test_checkdb_employee(self):
        c = self.session.query(Employee).count()
        self.assertEqual(c , 20)

However, I need to somehow pass self.engine to the decorator.

@loadfixtures(self.engine, ['employee.json']) does not work, because self is not available there.

I created this solution for this, but is there a simpler (more readable) solution?

    def test_checkdb_employee(self):
        @loadfixtures(['employee.json'])
        def run(engine):
            c = self.session.query(Employee).count()
            self.assertEqual(c , 20)
        return run(self.engine)

     
#decorator
import pandas as pd
 
def loadfixtures( files):
    def decorator(func):
        def wrapped_func(*args, **kwargs):
            for file in files: 
                df = pd.read_json(Path(datafolder).joinpath(file))
                df.to_sql(file.split(".")[0], con=args[0],index=False, if_exists='append')
            return func(*args, **kwargs)
        return wrapped_func
    return decorator

(The engine is used in con=args[0])

Upvotes: 0

Views: 79

Answers (1)

Glenn D.J.
Glenn D.J.

Reputation: 1965

self is actually passed as the first argument to test_checkdb_employee. So to access your engine in the decorator to use as argument for the con parameter you could just do args[0].engine with args[0] being self (the TestCompany instance).

def loadfixtures(files):
    def decorator(func):
        def wrapped_func(*args, **kwargs):
            for file in files: 
                df = pd.read_json(Path(datafolder).joinpath(file))
                df.to_sql(file.split(".")[0], con=args[0].engine, index=False, if_exists='append')
            return func(*args, **kwargs)
        return wrapped_func
    return decorator

Upvotes: 1

Related Questions