Changed so that it works for both Rails 3 and Rails 4