Adding null checks in a few places that could result in NPE