Clean up and made Unit and E2E tests passing