Cleaned up code.
Minor change to test code.