Added assertEqual helper and minor refactoring