LoadTestingService.cs•15 kB
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
using Azure;
using Azure.Core;
using Azure.Developer.LoadTesting;
using Azure.ResourceManager;
using Azure.ResourceManager.LoadTesting;
using Azure.ResourceManager.Resources;
using AzureMcp.Core.Options;
using AzureMcp.Core.Services.Azure;
using AzureMcp.Core.Services.Azure.Subscription;
using AzureMcp.LoadTesting.Commands;
using AzureMcp.LoadTesting.Models.LoadTest;
using AzureMcp.LoadTesting.Models.LoadTestResource;
using AzureMcp.LoadTesting.Models.LoadTestRun;
namespace AzureMcp.LoadTesting.Services;
public class LoadTestingService(ISubscriptionService subscriptionService) : BaseAzureService, ILoadTestingService
{
    ISubscriptionService _subscriptionService = subscriptionService;
    public async Task<List<TestResource>> GetLoadTestResourcesAsync(string subscription, string? resourceGroup = null, string? testResourceName = null, string? tenant = null, RetryPolicyOptions? retryPolicy = null)
    {
        ValidateRequiredParameters(subscription);
        var subscriptionId = (await _subscriptionService.GetSubscription(subscription, tenant, retryPolicy)).Data.SubscriptionId;
        var credential = await GetCredential(tenant);
        var client = new ArmClient(credential);
        if (!string.IsNullOrEmpty(testResourceName))
        {
            var resourceId = LoadTestingResource.CreateResourceIdentifier(subscriptionId, resourceGroup, testResourceName);
            var response = await client.GetLoadTestingResource(resourceId).GetAsync();
            if (response == null)
            {
                throw new Exception($"Failed to retrieve Azure Load Testing resources: {response}");
            }
            return new List<TestResource>
            {
                new TestResource
                {
                    Id = response.Value.Data.Id!,
                    Name = response.Value.Data.Name,
                    Location = response.Value.Data.Location,
                    DataPlaneUri = response.Value.Data.DataPlaneUri,
                    ProvisioningState = response.Value.Data.ProvisioningState?.ToString(),
                }
            };
        }
        else
        {
            var rgResource = ResourceGroupResource.CreateResourceIdentifier(subscriptionId, resourceGroup);
            var response = client.GetResourceGroupResource(rgResource).GetLoadTestingResources().ToList();
            if (response == null || response.Count == 0)
            {
                throw new Exception($"Failed to retrieve Azure Load Testing resources: {response}");
            }
            var loadTestResources = new List<TestResource>();
            foreach (var resource in response)
            {
                loadTestResources.Add(new TestResource
                {
                    Id = resource.Data.Id!,
                    Name = resource.Data.Name,
                    Location = resource.Data.Location,
                    DataPlaneUri = resource.Data.DataPlaneUri,
                    ProvisioningState = resource.Data.ProvisioningState?.ToString(),
                });
            }
            return loadTestResources;
        }
    }
    public async Task<TestResource> CreateOrUpdateLoadTestingResourceAsync(string subscription, string resourceGroup, string? testResourceName = null, string? tenant = null, RetryPolicyOptions? retryPolicy = null)
    {
        ValidateRequiredParameters(subscription, resourceGroup);
        var subscriptionId = (await _subscriptionService.GetSubscription(subscription, tenant, retryPolicy)).Data.SubscriptionId;
        var credential = await GetCredential(tenant);
        var client = new ArmClient(credential);
        var rgResource = client.GetResourceGroupResource(ResourceGroupResource.CreateResourceIdentifier(subscriptionId, resourceGroup));
        if (testResourceName == null)
        {
            testResourceName = $"estRun_{DateTime.UtcNow:dd-MM-yyyy_HH:mm:ss tt}";
        }
        var location = rgResource.Get().Value.Data.Location;
        var response = await rgResource.GetLoadTestingResources().CreateOrUpdateAsync(WaitUntil.Completed, testResourceName, new LoadTestingResourceData(location));
        if (response == null || response.Value == null)
        {
            throw new Exception($"Failed to create or update Azure Load Testing resource: {response}");
        }
        return new TestResource
        {
            Id = response.Value.Data.Id!,
            Name = response.Value.Data.Name,
            Location = response.Value.Data.Location,
            DataPlaneUri = response.Value.Data.DataPlaneUri,
            ProvisioningState = response.Value.Data.ProvisioningState?.ToString(),
        };
    }
    public async Task<TestRun> GetLoadTestRunAsync(string subscription, string testResourceName, string testRunId, string? resourceGroup = null, string? tenant = null, RetryPolicyOptions? retryPolicy = null)
    {
        ValidateRequiredParameters(subscription, testResourceName, testRunId);
        var subscriptionId = (await _subscriptionService.GetSubscription(subscription, tenant, retryPolicy)).Data.SubscriptionId;
        var loadTestResource = await GetLoadTestResourcesAsync(subscriptionId, resourceGroup, testResourceName, tenant, retryPolicy);
        if (loadTestResource == null)
        {
            throw new Exception($"Load Test '{testResourceName}' not found in subscription '{subscriptionId}' and resource group '{resourceGroup}'.");
        }
        var dataPlaneUri = loadTestResource[0]?.DataPlaneUri;
        if (string.IsNullOrEmpty(dataPlaneUri))
        {
            throw new Exception($"Data Plane URI for Load Test '{testResourceName}' is not available.");
        }
        var credential = await GetCredential(tenant);
        var loadTestClient = new LoadTestRunClient(new Uri($"https://{dataPlaneUri}"), credential);
        var loadTestRunResponse = await loadTestClient.GetTestRunAsync(testRunId);
        if (loadTestRunResponse == null || loadTestRunResponse.IsError)
        {
            throw new Exception($"Failed to retrieve Azure Load Test Run: {loadTestRunResponse}");
        }
        var loadTestRun = loadTestRunResponse.Content;
        return JsonSerializer.Deserialize(loadTestRun, LoadTestJsonContext.Default.TestRun) ?? new TestRun();
    }
    public async Task<List<TestRun>> GetLoadTestRunsFromTestIdAsync(string subscription, string testResourceName, string testId, string? resourceGroup = null, string? tenant = null, RetryPolicyOptions? retryPolicy = null)
    {
        ValidateRequiredParameters(subscription, testResourceName, testId);
        var subscriptionId = (await _subscriptionService.GetSubscription(subscription, tenant, retryPolicy)).Data.SubscriptionId;
        var loadTestResource = await GetLoadTestResourcesAsync(subscriptionId, resourceGroup, testResourceName, tenant, retryPolicy);
        if (loadTestResource == null)
        {
            throw new Exception($"Load Test '{testResourceName}' not found in subscription '{subscriptionId}' and resource group '{resourceGroup}'.");
        }
        var dataPlaneUri = loadTestResource[0]?.DataPlaneUri;
        if (string.IsNullOrEmpty(dataPlaneUri))
        {
            throw new Exception($"Data Plane URI for Load Test '{testResourceName}' is not available.");
        }
        var credential = await GetCredential(tenant);
        var loadTestClient = new LoadTestRunClient(new Uri($"https://{dataPlaneUri}"), credential);
        var loadTestRunResponse = loadTestClient.GetTestRunsAsync(testId: testId);
        if (loadTestRunResponse == null)
        {
            throw new Exception($"Failed to retrieve Azure Load Test Run: {loadTestRunResponse}");
        }
        var testRuns = new List<TestRun>();
        await foreach (var binaryData in loadTestRunResponse)
        {
            var testRun = JsonSerializer.Deserialize<TestRun>(binaryData.ToString(), LoadTestJsonContext.Default.TestRun);
            if (testRun != null)
            {
                testRuns.Add(testRun);
            }
        }
        if (testRuns.Count == 0)
        {
            throw new Exception($"No test runs found for test ID '{testId}' in Load Test '{testResourceName}'.");
        }
        return testRuns;
    }
    public async Task<TestRun> CreateOrUpdateLoadTestRunAsync(string subscription, string testResourceName, string testId, string? testRunId = null, string? oldTestRunId = null, string? resourceGroup = null, string? tenant = null, string? displayName = null, string? description = null, bool? debugMode = false, RetryPolicyOptions? retryPolicy = null)
    {
        ValidateRequiredParameters(subscription, testResourceName, testRunId);
        var subscriptionId = (await _subscriptionService.GetSubscription(subscription, tenant, retryPolicy)).Data.SubscriptionId;
        var loadTestResource = await GetLoadTestResourcesAsync(subscriptionId, resourceGroup, testResourceName, tenant, retryPolicy);
        if (loadTestResource == null)
        {
            throw new Exception($"Load Test '{testResourceName}' not found in subscription '{subscriptionId}' and resource group '{resourceGroup}'.");
        }
        var dataPlaneUri = loadTestResource[0]?.DataPlaneUri;
        if (string.IsNullOrEmpty(dataPlaneUri))
        {
            throw new Exception($"Data Plane URI for Load Test '{testResourceName}' is not available.");
        }
        var credential = await GetCredential(tenant);
        var loadTestClient = new LoadTestRunClient(new Uri($"https://{dataPlaneUri}"), credential);
        TestRunRequest requestBody = new TestRunRequest
        {
            TestId = testId,
            DisplayName = displayName ?? $"TestRun_{DateTime.UtcNow:dd-MM-yyyy_HH:mm:ss tt}",
            Description = description,
            DebugLogsEnabled = debugMode ?? false,
            RequestDataLevel = debugMode == true ? RequestDataLevel.ERRORS : RequestDataLevel.NONE,
        };
        var loadTestRunResponse = await loadTestClient.BeginTestRunAsync(0, testRunId, RequestContent.Create(JsonSerializer.Serialize(requestBody, LoadTestJsonContext.Default.TestRunRequest)), oldTestRunId: oldTestRunId);
        if (loadTestRunResponse == null)
        {
            throw new Exception($"Failed to retrieve Azure Load Test Run: {loadTestRunResponse}");
        }
        var loadTestRun = loadTestRunResponse.WaitForCompletionAsync().Result.Value.ToString();
        return JsonSerializer.Deserialize<TestRun>(loadTestRun, LoadTestJsonContext.Default.TestRun) ?? new TestRun();
    }
    public async Task<Test> GetTestAsync(string subscription, string testResourceName, string testId, string? resourceGroup = null, string? tenant = null, RetryPolicyOptions? retryPolicy = null)
    {
        ValidateRequiredParameters(subscription, testResourceName, testId);
        var subscriptionId = (await _subscriptionService.GetSubscription(subscription, tenant, retryPolicy)).Data.SubscriptionId;
        var loadTestResource = await GetLoadTestResourcesAsync(subscriptionId, resourceGroup, testResourceName, tenant, retryPolicy);
        if (loadTestResource == null)
        {
            throw new Exception($"Load Test '{testResourceName}' not found in subscription '{subscriptionId}' and resource group '{resourceGroup}'.");
        }
        var dataPlaneUri = loadTestResource[0]?.DataPlaneUri;
        if (string.IsNullOrEmpty(dataPlaneUri))
        {
            throw new Exception($"Data Plane URI for Load Test '{testResourceName}' is not available.");
        }
        var credential = await GetCredential(tenant);
        var loadTestClient = new LoadTestAdministrationClient(new Uri($"https://{dataPlaneUri}"), credential);
        var loadTestResponse = await loadTestClient.GetTestAsync(testId);
        if (loadTestResponse == null || loadTestResponse.IsError)
        {
            throw new Exception($"Failed to retrieve Azure Load Test: {loadTestResponse}");
        }
        var loadTest = loadTestResponse.Content.ToString();
        return JsonSerializer.Deserialize(loadTest, LoadTestJsonContext.Default.Test) ?? new Test();
    }
    public async Task<Test> CreateTestAsync(string subscription, string testResourceName, string testId, string? resourceGroup = null,
        string? displayName = null, string? description = null,
        int? duration = 20, int? virtualUsers = 50, int? rampUpTime = 1, string? endpointUrl = null, string? tenant = null, RetryPolicyOptions? retryPolicy = null)
    {
        ValidateRequiredParameters(subscription, testResourceName, testId);
        var subscriptionId = (await _subscriptionService.GetSubscription(subscription, tenant, retryPolicy)).Data.SubscriptionId;
        var loadTestResource = await GetLoadTestResourcesAsync(subscriptionId, resourceGroup, testResourceName, tenant, retryPolicy);
        if (loadTestResource == null)
        {
            throw new Exception($"Load Test '{testResourceName}' not found in subscription '{subscriptionId}' and resource group '{resourceGroup}'.");
        }
        var dataPlaneUri = loadTestResource[0]?.DataPlaneUri;
        if (string.IsNullOrEmpty(dataPlaneUri))
        {
            throw new Exception($"Data Plane URI for Load Test '{testResourceName}' is not available.");
        }
        var credential = await GetCredential(tenant);
        var loadTestClient = new LoadTestAdministrationClient(new Uri($"https://{dataPlaneUri}"), credential);
        OptionalLoadTestConfig optionalLoadTestConfig = new OptionalLoadTestConfig
        {
            Duration = (duration ?? 20) * 60, // Convert minutes to seconds
            EndpointUrl = endpointUrl ?? "https://example.com",
            RampUpTime = (rampUpTime ?? 1) * 60, // Convert minutes to seconds
            VirtualUsers = virtualUsers ?? 50,
        };
        TestRequestPayload testRequestPayload = new TestRequestPayload
        {
            TestId = testId,
            DisplayName = displayName ?? "Test_" + DateTime.UtcNow.ToString("dd/MM/yyyy") + "_" + DateTime.UtcNow.ToString("HH:mm:ss"),
            Description = description,
            LoadTestConfiguration = new LoadTestConfiguration
            {
                OptionalLoadTestConfig = optionalLoadTestConfig,
                QuickStartTest = true, // Set to true for quick start tests (URL BASIC Test)
            }
        };
        var loadTestResponse = await loadTestClient.CreateOrUpdateTestAsync(testId, RequestContent.Create(JsonSerializer.Serialize(testRequestPayload, LoadTestJsonContext.Default.TestRequestPayload)));
        if (loadTestResponse == null || loadTestResponse.IsError)
        {
            throw new Exception($"Failed to create Azure Load Test: {loadTestResponse}");
        }
        var loadTest = loadTestResponse.Content.ToString();
        return JsonSerializer.Deserialize(loadTest, LoadTestJsonContext.Default.Test) ?? new Test();
    }
}