Test your AI Agents with the all-new Agent to Agent Testing Platform.Learn More

How to use overwrite_user method in locust

Best Python code snippet using locust

create_user.py

Source:create_user.py

1#!/usr/bin/env python2import re3import sys4from passlib.context import CryptContext5from sqlalchemy import create_engine6from sqlalchemy.exc import IntegrityError7from sqlalchemy.orm.session import sessionmaker8from pylib.base.flags import Flags9from log import LOG10from models.alchemy.user import User, UserRoles, UserStatusEnum11from models.alchemy.query_policy import QueryPolicyRole12# I have to import this model since it is referenced by the `User` model13# pylint:disable=W061114from models.alchemy.security_group import GroupUsers15from models.alchemy.permission import Role16from models.alchemy.dashboard import Dashboard17from util.credentials.generate import generate_secure_password18from util.credentials.provider import CredentialProvider19from web.server.data.data_access import Transaction20from web.server.configuration.instance import load_instance_configuration_from_file21PASSWORD_ENCRYPTION_SCHEME = ['bcrypt']22PERMISSIVE_EMAIL_REGEX = re.compile(r'[^@]+@[^@]+\.[^@]+')23def get_user_string(user):24    return ('%s - (%s, %s)') % (user.username, user.first_name, user.last_name)25def is_email_address(username):26    return PERMISSIVE_EMAIL_REGEX.match(username) != None27def hash_password(password):28    # BCrypt has a built-in salt29    # https://stackoverflow.com/questions/6832445/how-can-bcrypt-have-built-in-salts30    crypt_context = CryptContext(schemes=PASSWORD_ENCRYPTION_SCHEME)31    return crypt_context.encrypt(password)32def create_user(33    transaction,34    username,35    first_name,36    last_name,37    plaintext_password=None,38    is_site_admin=False,39    overwrite_user=False,40    status=UserStatusEnum.ACTIVE,41):42    if not plaintext_password:43        plaintext_password = generate_secure_password()44    hashed_password = hash_password(plaintext_password)45    existing_user = transaction.find_one_by_fields(User, False, {'username': username})46    if existing_user and overwrite_user:47        LOG.info('User \'%s\' already exists but will be overwritten.', username)48    elif existing_user:49        message = (50            'User with username \'%s\' already exists. Specify \'-o\' to ovewrite.'51            % username52        )53        LOG.error(message)54        raise ValueError(message)55    if existing_user:56        if first_name:57            existing_user.first_name = first_name58        if last_name:59            existing_user.last_name = last_name60        existing_user.username = username61        existing_user.password = hashed_password62        existing_user.status_id = status.value63    else:64        existing_user = User(65            first_name=first_name,66            last_name=last_name,67            username=username,68            password=hashed_password,69            status_id=status.value,70        )71    new_user = transaction.add_or_update(existing_user, flush=True)72    if is_site_admin:73        site_admin_role = transaction.find_one_by_fields(Role, False, {'name': 'admin'})74        if not site_admin_role:75            message = (76                'Unable to find the site administrator role. It may have been deleted. '77                'Contact an Engineer for assistance. '78            )79            LOG.info(message)80            raise ValueError(message)81        user_role = transaction.find_one_by_fields(82            UserRoles, True, {'user_id': new_user.id, 'role_id': site_admin_role.id}83        )84        if user_role:85            LOG.info(86                '\'%s\' is already a site-administrator. ', get_user_string(new_user)87            )88        else:89            site_admin_user_role = UserRoles(90                user_id=new_user.id, role_id=site_admin_role.id91            )92            transaction.add_or_update(site_admin_user_role, flush=True)93            LOG.info(94                'Marked \'%s\' as a site-administrator. ', get_user_string(new_user)95            )96    return (new_user, plaintext_password)97def main():98    Flags.PARSER.add_argument(99        '-d',100        '--sql_connection_string',101        type=str,102        required=False,103        help='The SQL Connection String to use to connect to the SQL '104        'Database. Can also be specified via the \'DATABASE_URL\' '105        'environment variable. The inline parameter takes priority'106        'over the environment variable.',107    )108    Flags.PARSER.add_argument(109        '-u',110        '--username',111        type=str,112        required=False,113        help='The username of the user. MUST be a Zenysis e-mail address.',114    )115    Flags.PARSER.add_argument(116        '-f', '--first_name', type=str, required=False, help='The user\'s first name.'117    )118    Flags.PARSER.add_argument(119        '-l', '--last_name', type=str, required=False, help='The user\'s last name. '120    )121    Flags.PARSER.add_argument(122        '-p',123        '--password',124        type=str,125        required=False,126        help='The user\'s password. If none specified, this will be '127        'auto-generated. ',128    )129    Flags.PARSER.add_argument(130        '-s',131        '--status',132        type=str,133        action='store',134        required=False,135        choices=[e.name for e in UserStatusEnum],136        default=UserStatusEnum.ACTIVE.name,137        help=(138            'The type of SSL configuration to use. '139            '1. ACTIVE - The will be able to login immediately. '140            '2. INACTIVE - The user will not be able to login unless an '141            'Administrator logs in and marks the user as active. '142            '3. PENDING - The user will not be able to login unless an '143            'Administrator logs in and sends the user an invite email. '144        ),145    )146    Flags.PARSER.add_argument(147        '-a',148        '--site_admin',149        action='store_true',150        required=False,151        default=False,152        help='If specified, make user an admin.',153    )154    Flags.PARSER.add_argument(155        '-o',156        '--overwrite',157        action='store_true',158        required=False,159        default=False,160        help='Overwrite the user if the specified username already exists.',161    )162    Flags.InitArgs()163    sql_connection_string = Flags.ARGS.sql_connection_string164    if not sql_connection_string:165        instance_configuration = load_instance_configuration_from_file()166        with CredentialProvider(instance_configuration) as credential_provider:167            sql_connection_string = credential_provider.get('SQLALCHEMY_DATABASE_URI')168    username = Flags.ARGS.username169    first_name = Flags.ARGS.first_name or None170    last_name = Flags.ARGS.last_name or None171    plaintext_password = Flags.ARGS.password172    is_site_admin = Flags.ARGS.site_admin173    # pylint: disable=E1136174    # The types defined in Flags match exactly those defined in the Enum175    # there will not be a key error176    status = UserStatusEnum[Flags.ARGS.status]177    overwrite_user = Flags.ARGS.overwrite178    if not username:179        LOG.error('You must provide a username.')180        return 5181    if not overwrite_user and (not first_name or not last_name):182        LOG.error(183            'You must provide a first and last name if you are creating a new user.'184        )185        return 2186    username = username.strip()187    first_name = first_name.strip() if first_name else None188    last_name = last_name.strip() if last_name else None189    if not is_email_address(username):190        LOG.error(191            'Username \'%s\' is not valid. It must be an e-mail address.', username192        )193        return 3194    Session = sessionmaker()195    engine = create_engine(sql_connection_string)196    Session.configure(bind=engine)197    session = Session()198    with Transaction(should_commit=None, get_session=lambda: session) as transaction:199        (new_user, plaintext_password) = create_user(200            transaction,201            username,202            first_name,203            last_name,204            plaintext_password,205            is_site_admin,206            overwrite_user,207            status,208        )209        LOG.info(210            'Successfully created/updated User \'%s\' with status \'%s\' and password \'%s\'.',211            get_user_string(new_user),212            status.name,213            plaintext_password,214        )215    return 0216if __name__ == '__main__':...

client.py

Source:client.py

...22        return self.get(url)23    def create_user(self, first_name, last_name):24        data = {"firstName": first_name, "lastName": last_name}25        return self.post(Urls.users, data=data)26    def overwrite_user(self, user_id, first_name, last_name):27        data = {"firstName": first_name, "lastName": last_name}28        url = Urls.user.format(id=user_id)29        return self.put(url, data=data)30    def update_user(self, user_id, first_name=None, last_name=None):31        data = {}32        if first_name:33            data["firstName"] = first_name34        if last_name:35            data["lastName"] = last_name36        url = Urls.user.format(id=user_id)37        return self.patch(url, data=data)38    def delete_user(self, user_id):39        url = Urls.user.format(id=user_id)40        return self.delete(url)41    @paginated(by_query_params=by_query_params_callable)42    def list_user_accounts_paginated(self, user_id):43        return self.get(Urls.accounts, params={"userId": user_id})44@dataclass45class User:46    user_id: int = json_field(json="userId")47    first_name: str = json_field(json="firstName")48    last_name: str = json_field(json="lastName")49@dataclass50class Account:51    account_name: str = json_field(json="accountName")52    number: str = json_field(json="number")53@dataclass54class AccountPage:55    results: List[Account] = json_field(json="results")56    page: int = json_field(json="page")57    next_page: Optional[int] = json_field(json="nextPage")58class ClientWithJson(Client):59    @unmarshal_response(List[User])60    def list_users(self):61        return super().list_users()62    @unmarshal_response(User)63    def get_user(self, user_id: int):64        return super().get_user(user_id)65    @unmarshal_response(User)66    def create_user(self, first_name, last_name):67        return super().create_user(first_name, last_name)68    @unmarshal_response(User)69    def overwrite_user(self, user_id, first_name, last_name):70        return super().overwrite_user(user_id, first_name, last_name)71    @unmarshal_response(User)72    def update_user(self, user_id, first_name=None, last_name=None):73        return super().update_user(user_id, first_name, last_name)74    def delete_user(self, user_id):75        return super().delete_user(user_id)76    @unmarshal_response(List[AccountPage])77    @paginated(by_query_params=by_query_params_callable)78    def list_user_accounts_paginated(self, user_id):...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.