Updated test to prevent NullPointerException