Updated code based on WordPress Tests