Refactor tests.
Update README
Clean up code