added logic to pick correct class name in messages