Resolve deprecation errors and adjust code to make testing easier Signed-off-by: Aayush Gupta <aayushgupta219@gmail.com>