updated tests to the new class structure