Clean up code in testing for production release