Add os import to grains tests