How to override methods using Knock with Rails 6

[rwehresmann]

Zeitwerk that comes with Rails 6 introduced some problems with the current version of Knock, and I could solve them with the work that's being done in the branch v2.2.

However, I have an application where I use Knock with Auth0 and I override a method for the purposes of the application:

initializers/knock.rb:

# Find our User by id in the test environment and auth0_uid in other environments.
Knock::AuthToken.class_eval do
  def entity_for(entity_class)
    key_to_find = Rails.env.test? ? :id : :auth0_uid
    entity_class.find_by(key_to_find => @payload['sub'])
  end
end

That isn't working anymore. I simply created a fork, rewritten this method there and I'm using it right now, no problems. But I would like to keep up with the current updates coming with the gem without managing my fork. Any ideas on how to accomplish this override now?