Cleaned up and added test function to Main