Integration testing in Umbraco 10+: working with examine

[Fact]
public async Task Get_ValidRequest_Returns200OKWithJsonObject()
{
    // arrange

    // act
    var response = await Client.GetAsync(_endpoint + "?skip=0&take=10");

    // assert
    response.StatusCode.Should().Be(HttpStatusCode.OK, because: "the request is valid");
    var result = await DeserializeResponseAsync<GetNewsResponse>(response);

    result.Should().NotBeNull(because: "this endpoint always returns an object");
    result!.Total.Should().BePositive(because: "results exist");
    result.Items
        .Should().HaveCountLessThanOrEqualTo(10, because: "up to 10 results have been requested")
        .And.HaveCountGreaterThan(0, because: "results exist");

    result.Items.Should().NotContain(e => e.Title == null || e.SubTitle == null || e.Id == 0);
}

public class ExamineWaitContext
{
    private TaskCompletionSource? _tcp = null;

    public Task Set()
    {
        _tcp?.TrySetCanceled();
        _tcp = new TaskCompletionSource();

        return _tcp.Task;
    }

    public void Fire()
    {
        _tcp?.TrySetResult();
    }
}

internal class CustomExamineIndexRebuilder : ExamineIndexRebuilder
{
    private readonly ExamineWaitContext _waitContext;

    public CustomExamineIndexRebuilder(IMainDom mainDom,
                                           IRuntimeState runtimeState,
                                           ILogger<ExamineIndexRebuilder> logger,
                                           IExamineManager examineManager,
                                           IEnumerable<IIndexPopulator> populators,
                                           IBackgroundTaskQueue backgroundTaskQueue,
                                           ExamineWaitContext waitContext)
        : base(mainDom, runtimeState, logger, examineManager, populators, backgroundTaskQueue)
    {
        _waitContext = waitContext;
    }

    public override void RebuildIndexes(bool onlyEmptyIndexes, TimeSpan? delay = null, bool useBackgroundThread = true)
    {
        // Force rebuilding of indexes on the current thread
        //    so that the integration tests can wait for the rebuild to complete.
        base.RebuildIndexes(onlyEmptyIndexes, delay, false);

        // After rebuilding, send a signal to the wait context that the rebuild is done
        //    This way, the integration tests know that rebuilding is finished and testing can begin
        _waitContext.Fire();
    }
}

private void ConfigureServices(IServiceCollection services)
{
    // First remove the original implementations
    services.RemoveAll<ExamineIndexRebuilder>();
    services.RemoveAll<IIndexRebuilder>();

    // Then add the new implementations
    services.AddSingleton<ExamineIndexRebuilder, CustomExamineIndexRebuilder>();
    services.AddSingleton<IIndexRebuilder, CustomExamineIndexRebuilder>();
    services.AddSingleton<ExamineWaitContext>();
}

public class CustomRebuildOnStartupHandlerState
{
    public bool _isReady;
    public bool _isReadSet;
    public object? _isReadyLock;
}

public sealed class CustomRebuildOnStartupHandler
    : INotificationHandler<UmbracoRequestBeginNotification>
{
    // The static private fields have been removed
    //    and replaced with the _state field,
    //    which is injected in the constructor
    private readonly ExamineIndexRebuilder _backgroundIndexRebuilder;
    private readonly IRuntimeState _runtimeState;
    private readonly CustomRebuildOnStartupHandlerState _state;
    private readonly ISyncBootStateAccessor _syncBootStateAccessor;

    public CustomRebuildOnStartupHandler(
        ISyncBootStateAccessor syncBootStateAccessor,
        ExamineIndexRebuilder backgroundIndexRebuilder,
        IRuntimeState runtimeState,
        CustomRebuildOnStartupHandlerState state)
    {
        _syncBootStateAccessor = syncBootStateAccessor;
        _backgroundIndexRebuilder = backgroundIndexRebuilder;
        _runtimeState = runtimeState;
        _state = state;
    }

    // This method should be a copy of the method in the original handler,
    //    except using the singleton state object instead of static fields
    public void Handle(UmbracoRequestBeginNotification notification)
    {
        if (_runtimeState.Level != RuntimeLevel.Run)
        {
            return;
        }

        // This method call now uses the fields in the state object,
        //    rather than the static fields
        LazyInitializer.EnsureInitialized(
            ref _state._isReady,
            ref _state._isReadSet,
            ref _state._isReadyLock,
            () =>
            {
                SyncBootState bootState = _syncBootStateAccessor.GetSyncBootState();

                // This method now passes TimeSpan.Zero to this method, so that we don't have to wait that 1 minute at the start.
                _backgroundIndexRebuilder.RebuildIndexes(
                    bootState != SyncBootState.ColdBoot,
                    TimeSpan.Zero);

                return true;
            });
    }
}

// The umbraco rebuild on startup handler gets removed, because the static fields inside of it break unit tests when running multiple at once.
//   If examine related tests start breaking, check if the implementation of this type has changed.
// The service is replaced with a custom notification handler which uses a singleton class instead of static fields
var sd = services.First(s => s.ImplementationType == typeof(RebuildOnStartupHandler));
services.Remove(sd);
services.Add(new UniqueServiceDescriptor(typeof(INotificationHandler<UmbracoRequestBeginNotification>), typeof(CustomRebuildOnStartupHandler), ServiceLifetime.Transient));
services.AddSingleton<CustomRebuildOnStartupHandlerState>();

public abstract class IntegrationTestBase
    : IDisposable
{
    // ... the usual fields
    private readonly Task _initializationTask;

    public IntegrationTestBase()
    {
        WebsiteFactory = new UmbracoApplicationFactory();
        _initializationTask = WebsiteFactory.StartAsync();

        var ssf = WebsiteFactory.Services.GetRequiredService<IServiceScopeFactory>();
        Scope = ssf.CreateAsyncScope();
        Client = CreateHttpClient();
    }

    protected Task PrepareAsync()
    {
        return _initializationTask;
    }

    // ... the rest of my integration test base class
}

public class UmbracoWebsiteFactory
    : WebApplicationFactory<Program>
{
    // ... The usual implementation of the website factory

    private bool _started = false;
    public async Task StartAsync()
    {
        if (!_started)
        {
            var waitContext = Services.GetRequiredService<ExamineWaitContext>();
            var waitHandle = waitContext.Set();

            await CreateClient().GetAsync("/");
            await waitHandle;
            _started = true;
        }
    }
}

// arrange
await PrepareAsync();