[#4] Fix test imports for new package structure