Update to tests. Fixing adding of classes