Add tests for classes