azure-search-openai-demo icon indicating copy to clipboard operation
azure-search-openai-demo copied to clipboard

Add Locust support for authentication

Open DuboisABB opened this issue 1 year ago • 2 comments

This issue is for a: (mark with an x)

- [ ] bug report -> please search issues before submitting
- [X] feature request
- [ ] documentation issue or request
- [ ] regression (a behavior that used to work and stopped in a new release)

Please modify locustfile.py to handle a web app where authentication is required (i.e. AZURE_USE_AUTHENTICATION=true and AZURE_ENABLE_UNAUTHENTICATED_ACCESS=false).

I attempted to use GitHub Copilot to add this feature but it doesn't work and I'm a bit over my head.

DuboisABB avatar Aug 25 '24 17:08 DuboisABB

I haven't got this working yet, but this is what I've attempted so far - sharing in case it helps. The token fetching isn't working for me yet unfortunately.

import json
import logging
import os
import random
import subprocess
import time

from azure.identity import AzureDeveloperCliCredential
from dotenv import load_dotenv
from locust import HttpUser, between, task


def load_azd_env():
    """Get path to current azd env file and load file using python-dotenv"""
    result = subprocess.run("azd env list -o json", shell=True, capture_output=True, text=True)
    if result.returncode != 0:
        raise Exception("Error loading azd env")
    env_json = json.loads(result.stdout)
    env_file_path = None
    for entry in env_json:
        if entry["IsDefault"]:
            env_file_path = entry["DotEnvPath"]
    if not env_file_path:
        raise Exception("No default azd env file found")
    logging.info(f"Loading azd env from {env_file_path}")
    load_dotenv(env_file_path, override=True)

class ChatUser(HttpUser):
    wait_time = between(5, 20)

    @task
    def ask_question(self):
        load_azd_env()
        token = AzureDeveloperCliCredential().get_token(f"api://{os.environ['AZURE_SERVER_APP_ID']}/access_as_user", tenant_id=os.getenv('AZURE_AUTH_TENANT_ID', os.getenv('AZURE_TENANT_ID')))

        self.client.get("/",
                        headers={"Authorization": f"Bearer {token.token}"})

pamelafox avatar Aug 28 '24 17:08 pamelafox

Thanks for tagging me, I'll try to take a look

mattgotteiner avatar Sep 12 '24 15:09 mattgotteiner