Added Unit class and fixed unit tests to correspond.