def METHOD_NAME(cls, company):
    """
    Set company for current user.
    Note that user.company_id would not be the current company
    but only the default company for the user.
    """
    context = {"allowed_company_ids": company.ids}
    if "allowed_company_ids" in cls.env.context:
        cls.env.context.pop("allowed_company_ids")
    cls.env.context = dict(**cls.env.context, **context)