Unit Test fixes