Update classes to new API updates