Adds package imports that I forgot in refactor