Updated tests and started writing tests for new features