Cleaned up testing code