Refactor to include OnComponentWillBeRemoved.

Had to remove some tests and added a new one. Might be there needs to be a couple of more tests added.