Get the user role on Login

@mathiasconradt, @konrad.sopala, i tried adding above into the rule like:

function (user, context, callback) {
  context.idToken['https://any-namespace/roles'] = context.authorization.roles;
  context.accessToken['https://any-namespace/roles'] = context.authorization.roles;
  return callback(null, user, context);
}

But it is giving me the following error when i am trying to test it:

ERROR: Cannot read property 'roles' of undefined

Is there some issue in this? Also, I saw a rule in the rule list to get the roles from database, i tried using that but did not get the roles in the response. Can you let me know if i am doing something wrong here? thank you.