(cleanup) Code clean up in test files