azure-search-openai-demo
azure-search-openai-demo copied to clipboard
Add Locust support for authentication
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.
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}"})
Thanks for tagging me, I'll try to take a look