Fixed broken tests and changes from previous classes