Files
UnrealEngine/Engine/Source/Programs/Horde/HordeAgent.Tests/Leases/LeaseManagerTests.cs
2025-05-18 13:04:45 +08:00

181 lines
6.4 KiB
C#

// Copyright Epic Games, Inc. All Rights Reserved.
using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using EpicGames.Core;
using EpicGames.Horde;
using EpicGames.Horde.Agents;
using EpicGames.Horde.Agents.Leases;
using EpicGames.Horde.Agents.Sessions;
using HordeAgent.Leases;
using HordeAgent.Services;
using HordeAgent.Tests.Services;
using HordeCommon.Rpc;
using HordeCommon.Rpc.Messages;
using HordeCommon.Rpc.Tasks;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.VisualStudio.TestTools.UnitTesting;
using OpenTelemetry.Trace;
namespace HordeAgent.Tests.Leases;
public class FakeCapabilitiesService : ICapabilitiesService
{
public RpcAgentCapabilities Capabilities { get; set; } = new ();
public Task<RpcAgentCapabilities> GetCapabilitiesAsync(DirectoryReference? workingDir)
{
return Task.FromResult(Capabilities);
}
}
public class SessionStub(AgentId agentId, SessionId sessionId, DirectoryReference workingDir, IHordeClient hordeClient) : ISession
{
public AgentId AgentId { get; } = agentId;
public SessionId SessionId { get; } = sessionId;
public DirectoryReference WorkingDir { get; } = workingDir;
public IHordeClient HordeClient { get; } = hordeClient;
public async ValueTask DisposeAsync()
{
if (HordeClient is IAsyncDisposable disposableClient)
{
await disposableClient.DisposeAsync();
}
GC.SuppressFinalize(this);
}
}
internal class TestHandlerFactory : LeaseHandlerFactory<TestTask>
{
private readonly LeaseResult _result = new((byte[]?)null);
public TestHandlerFactory(LeaseResult? leaseResult = null)
{
_result = leaseResult ?? _result;
}
public override LeaseHandler<TestTask> CreateHandler(RpcLease lease)
{
return new Handler(lease, _result);
}
private class Handler(RpcLease rpcLease, LeaseResult leaseResult) : LeaseHandler<TestTask>(rpcLease)
{
protected override Task<LeaseResult> ExecuteAsync(ISession session, LeaseId leaseId, TestTask message, Tracer tracer, ILogger logger, CancellationToken cancellationToken)
{
logger.LogInformation("Executed TestTask");
return Task.FromResult(leaseResult);
}
}
}
[TestClass]
public sealed class LeaseManagerTests
{
private readonly FakeHordeRpcServer _hordeServer;
private readonly LeaseManager _leaseManager;
public LeaseManagerTests()
{
_hordeServer = new FakeHordeRpcServer(CreateConsoleLogger<FakeHordeRpcServer>());
_leaseManager = CreateLeaseManager(_hordeServer.GetHordeClient());
}
[TestMethod]
public async Task Run_SingleLease_FinishesSuccessfullyAsync()
{
using CancellationTokenSource cts = new(5000);
_hordeServer.ScheduleTestLease();
Task<SessionResult> runTask = _leaseManager.RunAsync(cts.Token);
_leaseManager.OnLeaseFinished += (lease, result) =>
{
Assert.AreEqual(RpcAgentStatus.Ok, _hordeServer.LastReportedStatus);
_hordeServer.SetAgentStatus(RpcAgentStatus.Stopped);
};
Assert.AreEqual(new SessionResult(SessionOutcome.BackOff, SessionReason.Completed), await runTask);
}
[TestMethod]
public async Task UpdateSession_SendsCapabilitiesInFinalUpdate_WhenSessionTerminates_Async()
{
using CancellationTokenSource cts = new(5000);
LeaseResult leaseResult = new (new SessionResult((_, _) => Task.CompletedTask));
RpcAgentCapabilities fooCaps = new();
fooCaps.Properties.Add("foo");
FakeCapabilitiesService capsService = new () { Capabilities = fooCaps };
LeaseManager leaseManager = CreateLeaseManager(_hordeServer.GetHordeClient(), capsService, leaseResult);
_hordeServer.ScheduleTestLease();
SessionResult sessionResult = await leaseManager.RunAsync(cts.Token);
_hordeServer.UpdateSessionRequests.Writer.Complete();
List<RpcUpdateSessionRequest> requests = await _hordeServer.UpdateSessionRequests.Reader.ReadAllAsync(cts.Token).ToListAsync(cts.Token);
Assert.AreEqual(3, requests.Count);
Assert.IsTrue(requests[2].Capabilities.Properties.Contains("foo"));
Assert.AreEqual(new SessionResult(SessionOutcome.RunCallback, SessionReason.Completed), sessionResult);
}
[TestMethod]
public async Task TerminateGracefully_LeaseIsActive_Async()
{
using CancellationTokenSource cts = new(5000);
_leaseManager.TerminateSessionAfterLease = true;
_hordeServer.ScheduleTestLease();
Assert.AreEqual(new SessionResult(SessionOutcome.Terminate, SessionReason.Completed), await _leaseManager.RunAsync(cts.Token));
Assert.AreEqual(RpcAgentStatus.Stopped, _hordeServer.AgentStatus);
}
[TestMethod]
public async Task TerminateGracefully_NoLeaseActive_Async()
{
using CancellationTokenSource cts = new(5000);
_leaseManager.TerminateSessionAfterLease = true;
Assert.AreEqual(new SessionResult(SessionOutcome.Terminate, SessionReason.Completed), await _leaseManager.RunAsync(cts.Token));
Assert.AreEqual(RpcAgentStatus.Stopped, _hordeServer.AgentStatus);
}
/// <summary>
/// Create a console logger for tests
/// </summary>
/// <typeparam name="T">Type to instantiate</typeparam>
/// <returns>A logger</returns>
private static ILogger<T> CreateConsoleLogger<T>()
{
using ILoggerFactory loggerFactory = LoggerFactory.Create(builder =>
{
builder.SetMinimumLevel(LogLevel.Debug);
builder.AddSimpleConsole(options => { options.SingleLine = true; });
});
return loggerFactory.CreateLogger<T>();
}
[SuppressMessage("Reliability", "CA2000:Dispose objects before losing scope")]
private static LeaseManager CreateLeaseManager(IHordeClient hordeClient, FakeCapabilitiesService? capsService = null, LeaseResult? leaseResult = null)
{
DirectoryReference tempWorkingDir = new(Path.Join(Path.GetTempPath(), Path.GetRandomFileName()));
AgentSettings settings = new() { WriteStepOutputToLogger = true };
ISession session = new SessionStub(new AgentId("testAgent"), SessionId.Parse("aaaaaaaaaaaaaaaaaaaaaaaa"), tempWorkingDir, hordeClient);
TestOptionsMonitor<AgentSettings> settingsOptions = new (settings);
StatusService statusService = new(settingsOptions, NullLogger<StatusService>.Instance);
return new LeaseManager(
session,
capsService ?? new FakeCapabilitiesService(),
statusService,
new DefaultSystemMetrics(),
new List<LeaseHandlerFactory>() { new TestHandlerFactory(leaseResult) },
new LeaseLoggerFactory(settingsOptions, CreateConsoleLogger<LeaseLoggerFactory>()),
settingsOptions,
TracerProvider.Default.GetTracer("TestTracer"),
CreateConsoleLogger<LeaseManager>());
}
}