Tests refactoring to use traits instead of Story class