Just test code clean-up