Code cleanup and test update.