Update code to please Travis CI