Updated Unit Test setup