Indexed large data sources in steps so the others get their turn (#1052)

This commit is contained in:
Thorsten Sommer authored and GitHub committed 2026-10-11 20:48:40 +02:00
1 parent fe4a64f382
commit a3c6b29313
26 files changed
+1932 -334

No files matched your search

+313
View File
@@ -0,0 +1,313 @@
using AIStudio.Tools.Services;
namespace AIStudio.Tests.Tools;
/// <summary>
/// Checks which runs the queue of the indexing runs hands out, and why each of them happens.
/// </summary>
/// <remarks>
/// File watchers, intervals and clicks ask for runs, often many at once for the same data source. A
/// data source which runs twice embeds everything twice, at a cloud embedding provider for real
/// money. One which never runs again goes out of date without anybody noticing. A run in the wrong
/// mode is the subtle case: only a manual retry signs in to a mailbox again after its server refused
/// a sign-in. The order counts as well: a mailbox which waits behind a large folder misses its new
/// mails for hours.
/// </remarks>
[TestFixture]
public sealed class DataSourceRunQueueTests
{
private const string DATA_SOURCE_ID = "6f1d6a4e-6a5e-4c62-9a4f-0f2d2c8b7a11";
private const string OTHER_DATA_SOURCE_ID = "5a1e7c3d-9b2f-4d8e-a6c4-3f0b9d2e1c75";
private const string MAILBOX_ID = "0c8e2b4f-3d7a-4e19-b5c6-8a2f1d9e7b30";
[Test]
public async Task RequestsForAWaitingDataSourceMergeIntoOneRun()
{
var queue = new DataSourceRunQueue();
var first = await Request(queue, DataSourceEmbeddingRefreshMode.WATCHER_HASH_CHECK);
var second = await Request(queue, DataSourceEmbeddingRefreshMode.HASH_CHECK);
Assert.Multiple(() =>
{
Assert.That(first, Is.EqualTo(DataSourceQueueRequestResult.QUEUED));
Assert.That(second, Is.EqualTo(DataSourceQueueRequestResult.ALREADY_QUEUED));
Assert.That(DequeueAll(queue), Is.EqualTo(new[] { new DataSourceRunRequest(DATA_SOURCE_ID, DataSourceEmbeddingRefreshMode.WATCHER_HASH_CHECK) }), "The first request decides why the run happens.");
});
}
[Test]
public async Task RequestsDuringARunLeaveExactlyOneFollowUpRun()
{
var queue = new DataSourceRunQueue();
await Request(queue, DataSourceEmbeddingRefreshMode.HASH_CHECK);
Assert.That(queue.TryDequeue(out _), Is.True);
var first = await Request(queue, DataSourceEmbeddingRefreshMode.WATCHER_HASH_CHECK);
var second = await Request(queue, DataSourceEmbeddingRefreshMode.CONTINUATION);
var requeue = queue.TryComplete(DATA_SOURCE_ID, true, out var pendingRefreshMode);
if (requeue)
await queue.EnqueueAsync(DATA_SOURCE_ID, DataSourceQueueLane.FRONT);
Assert.Multiple(() =>
{
Assert.That(first, Is.EqualTo(DataSourceQueueRequestResult.RUNNING_MARKED_PENDING));
Assert.That(second, Is.EqualTo(DataSourceQueueRequestResult.RUNNING));
Assert.That(requeue, Is.True);
Assert.That(pendingRefreshMode, Is.EqualTo(DataSourceEmbeddingRefreshMode.WATCHER_HASH_CHECK), "The first request decides why the follow-up run happens.");
Assert.That(DequeueAll(queue), Is.EqualTo(new[] { new DataSourceRunRequest(DATA_SOURCE_ID, DataSourceEmbeddingRefreshMode.WATCHER_HASH_CHECK) }), "The running run may have read the data source before the changes, so one more run is needed, and one is enough.");
});
}
[Test]
public async Task ARunWithoutRequestsLeavesNoFollowUpRun()
{
var queue = new DataSourceRunQueue();
await Request(queue, DataSourceEmbeddingRefreshMode.HASH_CHECK);
queue.TryDequeue(out _);
var requeue = queue.TryComplete(DATA_SOURCE_ID, true, out _);
var nextRequest = queue.TryReserve(DATA_SOURCE_ID, true, DataSourceEmbeddingRefreshMode.HASH_CHECK);
Assert.Multiple(() =>
{
Assert.That(requeue, Is.False);
Assert.That(nextRequest, Is.EqualTo(DataSourceQueueRequestResult.QUEUED), "The data source must not count as running after its run completed.");
});
}
[Test]
public async Task AFollowUpRunOfADataSourceWhichIsGoneIsDropped()
{
var queue = new DataSourceRunQueue();
await Request(queue, DataSourceEmbeddingRefreshMode.HASH_CHECK);
queue.TryDequeue(out _);
await Request(queue, DataSourceEmbeddingRefreshMode.WATCHER_HASH_CHECK);
var requeue = queue.TryComplete(DATA_SOURCE_ID, false, out _);
var nextRequest = queue.TryReserve(DATA_SOURCE_ID, true, DataSourceEmbeddingRefreshMode.HASH_CHECK);
Assert.Multiple(() =>
{
Assert.That(requeue, Is.False);
Assert.That(nextRequest, Is.EqualTo(DataSourceQueueRequestResult.QUEUED), "Neither the run nor its follow-up run may linger.");
});
}
[Test]
public async Task ADataSourceClearedWhileItWaitedLeavesNoRunBehind()
{
var queue = new DataSourceRunQueue();
await Request(queue, DataSourceEmbeddingRefreshMode.HASH_CHECK);
queue.Clear(DATA_SOURCE_ID);
Assert.That(queue.TryDequeue(out _), Is.False, "The data source was deleted while it waited. Its id is still in the channel, yet it must not run.");
}
[Test]
public async Task ARepairedDataSourceRunsOnceAndInTheModeOfTheRepair()
{
var queue = new DataSourceRunQueue();
await Request(queue, DataSourceEmbeddingRefreshMode.HASH_CHECK);
// What the repair does: forget the waiting run, then ask for a manual retry.
queue.Clear(DATA_SOURCE_ID);
await Request(queue, DataSourceEmbeddingRefreshMode.MANUAL_RETRY);
Assert.That(DequeueAll(queue), Is.EqualTo(new[] { new DataSourceRunRequest(DATA_SOURCE_ID, DataSourceEmbeddingRefreshMode.MANUAL_RETRY) }), "The id is in the channel twice now. Running both copies would embed the data source twice, and a run in the old mode would not sign in to a mailbox again.");
}
[Test]
public async Task AManualRetryReplacesTheModeOfAWaitingRun()
{
var queue = new DataSourceRunQueue();
await Request(queue, DataSourceEmbeddingRefreshMode.INTERVAL_CHECK);
var retry = await Request(queue, DataSourceEmbeddingRefreshMode.MANUAL_RETRY);
var later = await Request(queue, DataSourceEmbeddingRefreshMode.TOOL_REQUEST);
Assert.Multiple(() =>
{
Assert.That(retry, Is.EqualTo(DataSourceQueueRequestResult.ALREADY_QUEUED));
Assert.That(later, Is.EqualTo(DataSourceQueueRequestResult.ALREADY_QUEUED));
Assert.That(DequeueAll(queue), Is.EqualTo(new[] { new DataSourceRunRequest(DATA_SOURCE_ID, DataSourceEmbeddingRefreshMode.MANUAL_RETRY) }), "Only a manual retry signs in to a mailbox again after its server refused a sign-in, and no later request may take that back.");
});
}
[Test]
public async Task AManualRetryReplacesTheModeOfAFollowUpRun()
{
var queue = new DataSourceRunQueue();
await Request(queue, DataSourceEmbeddingRefreshMode.INTERVAL_CHECK);
queue.TryDequeue(out _);
await Request(queue, DataSourceEmbeddingRefreshMode.INTERVAL_CHECK);
await Request(queue, DataSourceEmbeddingRefreshMode.MANUAL_RETRY);
await Request(queue, DataSourceEmbeddingRefreshMode.TOOL_REQUEST);
queue.TryComplete(DATA_SOURCE_ID, true, out var pendingRefreshMode);
Assert.That(pendingRefreshMode, Is.EqualTo(DataSourceEmbeddingRefreshMode.MANUAL_RETRY), "Only a manual retry signs in to a mailbox again after its server refused a sign-in, and no later request may take that back.");
}
[Test]
public async Task ARunOfTheStartupCheckGetsExactlyOneFollowUpRun()
{
var queue = new DataSourceRunQueue();
var begun = queue.TryBeginRun(DATA_SOURCE_ID);
var first = await Request(queue, DataSourceEmbeddingRefreshMode.CONTINUATION);
var second = await Request(queue, DataSourceEmbeddingRefreshMode.WATCHER_HASH_CHECK);
var requeue = queue.TryComplete(DATA_SOURCE_ID, true, out var pendingRefreshMode);
if (requeue)
await queue.EnqueueAsync(DATA_SOURCE_ID, DataSourceQueueLane.BACK);
Assert.Multiple(() =>
{
Assert.That(begun, Is.True);
Assert.That(first, Is.EqualTo(DataSourceQueueRequestResult.RUNNING_MARKED_PENDING), "The run began outside the queue, yet it runs all the same.");
Assert.That(second, Is.EqualTo(DataSourceQueueRequestResult.RUNNING));
Assert.That(requeue, Is.True);
Assert.That(pendingRefreshMode, Is.EqualTo(DataSourceEmbeddingRefreshMode.CONTINUATION));
Assert.That(DequeueAll(queue), Is.EqualTo(new[] { new DataSourceRunRequest(DATA_SOURCE_ID, DataSourceEmbeddingRefreshMode.CONTINUATION) }), "A data source which did its share of the work during the startup check must carry on afterward, exactly once.");
});
}
[Test]
public async Task TheStartupCheckLeavesOutADataSourceWhichRunsOrWaits()
{
var queue = new DataSourceRunQueue();
await Request(queue, DataSourceEmbeddingRefreshMode.HASH_CHECK);
var running = queue.TryBeginRun(OTHER_DATA_SOURCE_ID);
Assert.Multiple(() =>
{
Assert.That(running, Is.True);
Assert.That(queue.TryBeginRun(DATA_SOURCE_ID), Is.False, "The data source waits in the queue, so it would run twice in a row.");
Assert.That(queue.TryBeginRun(OTHER_DATA_SOURCE_ID), Is.False, "The data source runs already, so it would run twice at once.");
});
}
[Test]
public async Task WaitingForTheNextRunSkipsWhatWasCleared()
{
var queue = new DataSourceRunQueue();
await Request(queue, DataSourceEmbeddingRefreshMode.HASH_CHECK);
queue.Clear(DATA_SOURCE_ID);
var nextRun = queue.DequeueAsync(CancellationToken.None).AsTask();
var waitedForTheNextRun = !nextRun.IsCompleted;
await Request(queue, DataSourceEmbeddingRefreshMode.HASH_CHECK, OTHER_DATA_SOURCE_ID);
var run = await nextRun.WaitAsync(TimeSpan.FromSeconds(10));
Assert.Multiple(() =>
{
Assert.That(waitedForTheNextRun, Is.True, "Nothing was left to run.");
Assert.That(run, Is.EqualTo(new DataSourceRunRequest(OTHER_DATA_SOURCE_ID, DataSourceEmbeddingRefreshMode.HASH_CHECK)));
});
}
[Test]
public async Task AMailboxWaitsForOneStepAtMost()
{
var queue = new DataSourceRunQueue();
await Request(queue, DataSourceEmbeddingRefreshMode.CONTINUATION, DATA_SOURCE_ID, DataSourceQueueLane.BACK);
await Request(queue, DataSourceEmbeddingRefreshMode.CONTINUATION, OTHER_DATA_SOURCE_ID, DataSourceQueueLane.BACK);
Assert.That(queue.TryDequeue(out var step), Is.True);
// The interval of the mailbox comes due while a large folder does its step:
await Request(queue, DataSourceEmbeddingRefreshMode.INTERVAL_CHECK, MAILBOX_ID);
await CarryOn(queue, step!.DataSourceId);
Assert.Multiple(() =>
{
Assert.That(step.DataSourceId, Is.EqualTo(DATA_SOURCE_ID));
Assert.That(DequeueAll(queue).Select(run => run.DataSourceId), Is.EqualTo(new[] { MAILBOX_ID, OTHER_DATA_SOURCE_ID, DATA_SOURCE_ID }), "The mailbox must not wait behind every step of every large data source, which would add up to hours.");
});
}
[Test]
public async Task DataSourcesWhichCarryOnTakeTurns()
{
var queue = new DataSourceRunQueue();
await Request(queue, DataSourceEmbeddingRefreshMode.CONTINUATION, DATA_SOURCE_ID, DataSourceQueueLane.BACK);
await Request(queue, DataSourceEmbeddingRefreshMode.CONTINUATION, OTHER_DATA_SOURCE_ID, DataSourceQueueLane.BACK);
var steps = new List<string>();
for (var step = 0; step < 4; step++)
{
Assert.That(queue.TryDequeue(out var runRequest), Is.True);
steps.Add(runRequest!.DataSourceId);
await CarryOn(queue, runRequest.DataSourceId);
}
Assert.That(steps, Is.EqualTo(new[] { DATA_SOURCE_ID, OTHER_DATA_SOURCE_ID, DATA_SOURCE_ID, OTHER_DATA_SOURCE_ID }), "A data source which takes hours must not hold up another one which takes hours as well.");
}
[Test]
public async Task ADataSourceRepairedWhileItCarriedOnRunsOnceAndAtTheFront()
{
var queue = new DataSourceRunQueue();
await Request(queue, DataSourceEmbeddingRefreshMode.CONTINUATION, DATA_SOURCE_ID, DataSourceQueueLane.BACK);
await Request(queue, DataSourceEmbeddingRefreshMode.CONTINUATION, OTHER_DATA_SOURCE_ID, DataSourceQueueLane.BACK);
// What the repair does: forget the waiting run, which also drops the pass, then ask for a manual retry.
queue.Clear(DATA_SOURCE_ID);
await Request(queue, DataSourceEmbeddingRefreshMode.MANUAL_RETRY);
Assert.That(DequeueAll(queue), Is.EqualTo(new[] { new DataSourceRunRequest(DATA_SOURCE_ID, DataSourceEmbeddingRefreshMode.MANUAL_RETRY), new DataSourceRunRequest(OTHER_DATA_SOURCE_ID, DataSourceEmbeddingRefreshMode.CONTINUATION) }), "The copy left behind in the back lane must not run the data source a second time.");
}
[Test]
public async Task ACopyLeftBehindInTheFrontLaneLetsNoLaterRunOvertake()
{
var queue = new DataSourceRunQueue();
await Request(queue, DataSourceEmbeddingRefreshMode.WATCHER_HASH_CHECK);
await Request(queue, DataSourceEmbeddingRefreshMode.CONTINUATION, OTHER_DATA_SOURCE_ID, DataSourceQueueLane.BACK);
// The repair of a data source which waits in the front lane leaves a copy of its id behind there:
queue.Clear(DATA_SOURCE_ID);
await Request(queue, DataSourceEmbeddingRefreshMode.MANUAL_RETRY);
Assert.That(queue.TryDequeue(out var repair), Is.True);
await CarryOn(queue, repair!.DataSourceId);
Assert.Multiple(() =>
{
Assert.That(repair, Is.EqualTo(new DataSourceRunRequest(DATA_SOURCE_ID, DataSourceEmbeddingRefreshMode.MANUAL_RETRY)));
Assert.That(DequeueAll(queue).Select(run => run.DataSourceId), Is.EqualTo(new[] { OTHER_DATA_SOURCE_ID, DATA_SOURCE_ID }), "The next step of the repaired data source carries on its pass, so it waits its turn in the back lane.");
});
}
/// <summary>
/// Asks for a run the way DataSourceEmbeddingService does: reserve a place, and enqueue the data source when it got one.
/// </summary>
private static async Task<DataSourceQueueRequestResult> Request(DataSourceRunQueue queue, DataSourceEmbeddingRefreshMode refreshMode, string dataSourceId = DATA_SOURCE_ID, DataSourceQueueLane lane = DataSourceQueueLane.FRONT)
{
var result = queue.TryReserve(dataSourceId, true, refreshMode);
if (result is DataSourceQueueRequestResult.QUEUED)
await queue.EnqueueAsync(dataSourceId, lane);
return result;
}
/// <summary>
/// Ends the step of a running data source the way DataSourceEmbeddingService does when the step left more to do.
/// </summary>
private static async Task CarryOn(DataSourceRunQueue queue, string dataSourceId)
{
await Request(queue, DataSourceEmbeddingRefreshMode.CONTINUATION, dataSourceId);
if (queue.TryComplete(dataSourceId, true, out _))
await queue.EnqueueAsync(dataSourceId, DataSourceQueueLane.BACK);
}
private static List<DataSourceRunRequest> DequeueAll(DataSourceRunQueue queue)
{
var runs = new List<DataSourceRunRequest>();
while (queue.TryDequeue(out var runRequest))
runs.Add(runRequest);
return runs;
}
}
+251 -5
View File
@@ -7,6 +7,7 @@ using AIStudio.Tools.Databases.VectorStore;
using AIStudio.Tools.Services;
using AIStudio.Tools.Services.Indexing;
using Microsoft.AspNetCore.Components;
using Microsoft.Extensions.Logging.Abstractions;
namespace AIStudio.Tests.Tools;
@@ -27,6 +28,8 @@ public sealed class DocumentRunProgressTests
private const string EMBEDDING_ID = "b0a4c4d2-1f3e-4f0a-8c9d-5a6b7c8d9e01";
private static readonly DateTimeOffset NOW = new(2026, 10, 11, 9, 0, 0, TimeSpan.Zero);
[Test]
public async Task AProviderFailureKeepsTheDocumentsSkippedForGoodInTheStatus()
{
@@ -100,15 +103,234 @@ public sealed class DocumentRunProgressTests
var (progress, statuses, manifest) = CreateProgress(totalDocuments: 10);
progress.RecordUnchanged(4);
await progress.PauseRunAsync("SOURCE-HASH", "test", CancellationToken.None);
var outcome = await progress.PauseRunAsync("SOURCE-HASH", "test", CancellationToken.None);
Assert.Multiple(() =>
{
Assert.That(outcome, Is.EqualTo(IndexedRunOutcome.MORE_TO_DO));
Assert.That(manifest.SourceHash, Is.EqualTo("SOURCE-HASH"), "Without the hash, a mailbox stays out of every search for the hours its first sync takes.");
Assert.That(statuses, Is.Empty, "The status said the run is over, while the next one is about to carry on.");
});
}
[Test]
public async Task APausedRunWithoutAHashLeavesTheStoredOneAsItIs()
{
var (progress, statuses, manifest) = CreateProgress(totalDocuments: 10);
manifest.SourceHash = "HASH-OF-THE-LAST-COMPLETE-RUN";
progress.RecordUnchanged(4);
var outcome = await progress.PauseRunAsync(null, "test", CancellationToken.None);
Assert.Multiple(() =>
{
Assert.That(outcome, Is.EqualTo(IndexedRunOutcome.MORE_TO_DO));
Assert.That(manifest.SourceHash, Is.EqualTo("HASH-OF-THE-LAST-COMPLETE-RUN"), "The hash says that a run got through the whole folder. A half-indexed folder would count as complete, and the next run would skip it.");
Assert.That(statuses, Is.Empty);
});
}
[Test]
public async Task APauseAfterTheProviderFailedForEveryDocumentEndsTheRun()
{
var (progress, statuses, manifest) = CreateProgress(totalDocuments: 10);
var providerDown = new ProviderRequestException(ProviderRequestFailureReason.PROVIDER_UNAVAILABLE, "The embedding provider cannot be reached.");
await progress.RecordDocumentFailureAsync(Document("/tmp/test-data/a.pdf"), providerDown, CancellationToken.None);
await progress.RecordDocumentFailureAsync(Document("/tmp/test-data/b.pdf"), providerDown, CancellationToken.None);
var outcome = await progress.PauseRunAsync(null, "test", CancellationToken.None);
Assert.Multiple(() =>
{
Assert.That(outcome, Is.EqualTo(IndexedRunOutcome.DONE), "The next run would fail all the same, and the one after it, without end.");
Assert.That(statuses[^1].State, Is.EqualTo(DataSourceEmbeddingState.FAILED), "The run is over, so its row has to say so, and why.");
Assert.That(statuses[^1].LastError, Is.EqualTo("The embedding provider cannot be reached."));
Assert.That(manifest.SourceHash, Is.Empty, "A run which stopped halfway did not get through the folder.");
});
}
[Test]
public async Task DocumentsSkippedForGoodDoNotKeepARunWithAFailingProviderGoing()
{
var (progress, _, _) = CreateProgress(totalDocuments: 10);
await progress.RecordDocumentFailureAsync(Document("/tmp/test-data/scan.pdf"), new FileExtractionException(FileExtractionErrorCode.NO_TEXT_EXTRACTED, "No text."), CancellationToken.None);
await progress.RecordDocumentFailureAsync(Document("/tmp/test-data/report.pdf"), new ProviderRequestException(ProviderRequestFailureReason.PROVIDER_UNAVAILABLE, "The embedding provider cannot be reached."), CancellationToken.None);
var outcome = await progress.PauseRunAsync(null, "test", CancellationToken.None);
Assert.That(outcome, Is.EqualTo(IndexedRunOutcome.DONE), "A scanned document never gets as far as the provider, so it says nothing about whether the provider works.");
}
[Test]
public async Task APauseAfterFailuresOfTheDocumentsThemselvesCarriesOn()
{
var (progress, statuses, _) = CreateProgress(totalDocuments: 10);
await progress.RecordDocumentFailureAsync(Document("/tmp/test-data/a.pdf"), new IOException("The file changed."), CancellationToken.None);
await progress.RecordDocumentFailureAsync(Document("/tmp/test-data/b.pdf"), new ProviderRequestException(ProviderRequestFailureReason.TOO_MANY_REQUESTS, "Too many requests."), CancellationToken.None);
var outcome = await progress.PauseRunAsync(null, "test", CancellationToken.None);
Assert.Multiple(() =>
{
Assert.That(outcome, Is.EqualTo(IndexedRunOutcome.MORE_TO_DO), "The pass skips both documents from now on, so the next step gets further. Ending here left every file behind a block of failing files unindexed.");
Assert.That(statuses.Select(status => status.State), Has.None.EqualTo(DataSourceEmbeddingState.FAILED));
});
}
[Test]
public async Task InputsWhichCouldNotBeLookedAtDoNotEndAPausedRun()
{
var (progress, statuses, _) = CreateProgress(totalDocuments: 10, failedInputs: 1);
// A step which spent its share on a single scanned document without a text layer:
await progress.RecordDocumentFailureAsync(Document("/tmp/test-data/scan.pdf"), new FileExtractionException(FileExtractionErrorCode.NO_TEXT_EXTRACTED, "No text."), CancellationToken.None);
var outcome = await progress.PauseRunAsync(null, "test", CancellationToken.None);
Assert.Multiple(() =>
{
Assert.That(outcome, Is.EqualTo(IndexedRunOutcome.MORE_TO_DO), "A locked subfolder fails in every run, however much the run gets done otherwise, and the scanned document is never read again. Ending here left the rest of the folder unindexed.");
Assert.That(statuses.Select(status => status.State), Has.None.EqualTo(DataSourceEmbeddingState.FAILED));
});
}
[Test]
public void TheShareCountsFromTheFirstQuestion()
{
var (progress, _, _) = CreateProgress(totalDocuments: 10, share: new RunShare(TimeSpan.FromMinutes(5)));
var first = progress.HasDoneItsShare(NOW);
var shortlyBefore = progress.HasDoneItsShare(NOW + TimeSpan.FromMinutes(5) - TimeSpan.FromSeconds(1));
var atTheEnd = progress.HasDoneItsShare(NOW + TimeSpan.FromMinutes(5));
Assert.Multiple(() =>
{
Assert.That(first, Is.False, "However long finding the documents took, every run reads at least one. Otherwise, a folder whose files take longer to list than the share would never get anywhere.");
Assert.That(shortlyBefore, Is.False);
Assert.That(atTheEnd, Is.True);
});
}
[Test]
public async Task TheShareEndsAfterItsDocumentsWhetherIndexingThemWorkedOrNot()
{
var (progress, _, _) = CreateProgress(totalDocuments: 10, share: new RunShare(RunShare.DEFAULT_DURATION, 3));
progress.HasDoneItsShare(NOW);
await progress.RecordDocumentIndexedAsync(Document("/tmp/test-data/a.pdf"), 4, true, CancellationToken.None);
await progress.RecordDocumentFailureAsync(Document("/tmp/test-data/b.pdf"), new IOException("The file changed."), CancellationToken.None);
var afterTwo = progress.HasDoneItsShare(NOW);
await progress.RecordDocumentFailureAsync(Document("/tmp/test-data/scan.pdf"), new FileExtractionException(FileExtractionErrorCode.NO_TEXT_EXTRACTED, "No text."), CancellationToken.None);
var afterThree = progress.HasDoneItsShare(NOW);
Assert.Multiple(() =>
{
Assert.That(afterTwo, Is.False);
Assert.That(afterThree, Is.True, "Reading a document costs its time, whether indexing it worked or not.");
});
}
[Test]
public void DocumentsWhichWereNotReadDoNotCountTowardsTheShare()
{
var (progress, _, _) = CreateProgress(totalDocuments: 1_000, share: new RunShare(RunShare.DEFAULT_DURATION, 3));
progress.RecordUnchanged(900);
progress.RecordStillUnreadable("/tmp/test-data/scan.pdf", new PermanentIndexingFailureRecord("AAAA", FileExtractionErrorCode.NO_TEXT_EXTRACTED, "No text could be read.", DateTimeOffset.UtcNow));
Assert.That(progress.HasDoneItsShare(NOW), Is.False, "Leaving a document alone costs next to nothing. Counting it would end the step of a mailbox after its first few mails.");
}
[Test]
public async Task AFailureOfAnEarlierStepKeepsItsPlaceInTheNextOne()
{
var pass = NewPass();
var (firstStep, _, _) = CreateProgress(totalDocuments: 10, pass: pass);
await firstStep.RecordDocumentFailureAsync(Document("/tmp/test-data/report.pdf"), new IOException("The file changed."), CancellationToken.None);
var (nextStep, statuses, _) = CreateProgress(totalDocuments: 10, pass: pass);
var skipped = nextStep.SkipIfFailedEarlierInPass("/tmp/test-data/report.pdf", "FINGERPRINT");
nextStep.Publish();
Assert.Multiple(() =>
{
Assert.That(skipped, Is.True, "Reading it once more costs time, with a cloud provider money as well, and ends the same way.");
Assert.That(statuses[^1].FailedDocuments, Is.EqualTo(1));
Assert.That(statuses[^1].Failures.Select(failure => failure.DocumentKey), Is.EquivalentTo(["/tmp/test-data/report.pdf"]), "The list lost the failures of every earlier step.");
Assert.That(statuses[^1].LastError, Is.Not.Empty);
});
}
[Test]
public async Task ADocumentWhichChangedSinceItFailedIsTriedAgain()
{
var pass = NewPass();
var (firstStep, _, _) = CreateProgress(totalDocuments: 10, pass: pass);
await firstStep.RecordDocumentFailureAsync(Document("/tmp/test-data/report.pdf"), new IOException("The file changed."), CancellationToken.None);
var (nextStep, _, _) = CreateProgress(totalDocuments: 10, pass: pass);
Assert.That(nextStep.SkipIfFailedEarlierInPass("/tmp/test-data/report.pdf", "CHANGED-FINGERPRINT"), Is.False, "The change may well be what makes it readable, e.g. the program which wrote it is done now.");
}
[Test]
public async Task ADocumentIsTriedTwiceDuringAPassAtMost()
{
var pass = NewPass();
var (firstStep, _, _) = CreateProgress(totalDocuments: 10, pass: pass);
await firstStep.RecordDocumentFailureAsync(Document("/tmp/test-data/app.log", "app.log", "FINGERPRINT-1"), new IOException("The file changed."), CancellationToken.None);
var (secondStep, _, _) = CreateProgress(totalDocuments: 10, pass: pass);
await secondStep.RecordDocumentFailureAsync(Document("/tmp/test-data/app.log", "app.log", "FINGERPRINT-2"), new IOException("The file changed."), CancellationToken.None);
var (thirdStep, _, _) = CreateProgress(totalDocuments: 10, pass: pass);
Assert.That(thirdStep.SkipIfFailedEarlierInPass("/tmp/test-data/app.log", "FINGERPRINT-3"), Is.True, "A file which changes faster than it can be read would be read in every step, and the pass would never end.");
}
[Test]
public async Task ADocumentWhichWasIndexedAfterAllIsNoLongerSkipped()
{
var pass = NewPass();
var (firstStep, _, _) = CreateProgress(totalDocuments: 10, pass: pass);
await firstStep.RecordDocumentFailureAsync(Document("/tmp/test-data/report.pdf"), new IOException("The file changed."), CancellationToken.None);
await firstStep.RecordDocumentIndexedAsync(Document("/tmp/test-data/report.pdf"), 4, false, CancellationToken.None);
var (nextStep, _, _) = CreateProgress(totalDocuments: 10, pass: pass);
Assert.That(nextStep.SkipIfFailedEarlierInPass("/tmp/test-data/report.pdf", "FINGERPRINT"), Is.False, "An indexed document showed up as failed for the rest of the pass.");
}
[Test]
public async Task AProviderFailureIsReportedOncePerPass()
{
var message = $"Too many requests. Test {Guid.NewGuid()}";
var collector = new ErrorCollector(message);
MessageBus.INSTANCE.RegisterComponent(collector);
MessageBus.INSTANCE.ApplyFilters(collector, [], [Event.SHOW_ERROR]);
try
{
var pass = NewPass();
var (firstStep, _, _) = CreateProgress(totalDocuments: 10, pass: pass);
await firstStep.RecordDocumentFailureAsync(Document("/tmp/test-data/a.pdf"), new ProviderRequestException(ProviderRequestFailureReason.TOO_MANY_REQUESTS, message), CancellationToken.None);
var (nextStep, _, _) = CreateProgress(totalDocuments: 10, pass: pass);
await nextStep.RecordDocumentFailureAsync(Document("/tmp/test-data/b.pdf"), new ProviderRequestException(ProviderRequestFailureReason.TOO_MANY_REQUESTS, message), CancellationToken.None);
var reportedDuringThePass = collector.Count;
var (nextPass, _, _) = CreateProgress(totalDocuments: 10, pass: NewPass());
await nextPass.RecordDocumentFailureAsync(Document("/tmp/test-data/a.pdf"), new ProviderRequestException(ProviderRequestFailureReason.TOO_MANY_REQUESTS, message), CancellationToken.None);
Assert.Multiple(() =>
{
Assert.That(reportedDuringThePass, Is.EqualTo(1), "A rate limit which lasts showed a message every five minutes.");
Assert.That(collector.Count, Is.EqualTo(2), "A new pass is a new attempt, which the user hears about should it fail again.");
});
}
finally
{
MessageBus.INSTANCE.Unregister(collector);
}
}
[Test]
public void EveryStatusOfARunCarriesTheLastSync()
{
@@ -178,7 +400,9 @@ public sealed class DocumentRunProgressTests
Assert.That(statuses[^1].Failures[0].DisplayName, Is.EqualTo("report.pdf"));
}
private static (DocumentRunProgress Progress, List<DataSourceEmbeddingStatus> Statuses, DataSourceEmbeddingManifest Manifest) CreateProgress(int totalDocuments, IIndexedDataSource? indexedDataSource = null)
private static IndexingPass NewPass() => new(new DataSourceRunSignature("TEST-EMBEDDING-SIGNATURE"));
private static (DocumentRunProgress Progress, List<DataSourceEmbeddingStatus> Statuses, DataSourceEmbeddingManifest Manifest) CreateProgress(int totalDocuments, IIndexedDataSource? indexedDataSource = null, RunShare? share = null, int failedInputs = 0, IndexingPass? pass = null)
{
var statuses = new List<DataSourceEmbeddingStatus>();
var manifest = new DataSourceEmbeddingManifest();
@@ -201,18 +425,20 @@ public sealed class DocumentRunProgressTests
new NoVectorStoreClient("Test vector store", "Not available in tests."),
new NoIndexStoreClient("Test index store", "Not available in tests."),
manifest,
share ?? new RunShare(RunShare.DEFAULT_DURATION),
pass ?? NewPass(),
new SettingsManager(NullLogger<SettingsManager>.Instance, null!),
statuses.Add,
NullLogger.Instance);
return (new DocumentRunProgress(context, totalDocuments, 0, string.Empty, [], NullLogger.Instance), statuses, manifest);
return (new DocumentRunProgress(context, totalDocuments, failedInputs, string.Empty, [], NullLogger.Instance), statuses, manifest);
}
private static EmbeddingDocument Document(string path) => Document(path, Path.GetFileName(path));
private static EmbeddingDocument Document(string key, string displayName) => new(
private static EmbeddingDocument Document(string key, string displayName, string fingerprint = "FINGERPRINT") => new(
key,
new EmbeddingStateFile(IndexedDocumentIds.CreateParentId(DATA_SOURCE_ID, key), key, displayName, displayName, "pdf", "FINGERPRINT", 1024, DateTimeOffset.UnixEpoch, DateTimeOffset.UnixEpoch, DateTimeOffset.UnixEpoch, 0),
new EmbeddingStateFile(IndexedDocumentIds.CreateParentId(DATA_SOURCE_ID, key), key, displayName, displayName, "pdf", fingerprint, 1024, DateTimeOffset.UnixEpoch, DateTimeOffset.UnixEpoch, DateTimeOffset.UnixEpoch, 0),
displayName,
_ => NoChunks());
@@ -221,4 +447,24 @@ public sealed class DocumentRunProgressTests
await Task.CompletedTask;
yield break;
}
/// <summary>
/// Counts the error messages with one text, out of whatever other tests send meanwhile.
/// </summary>
private sealed class ErrorCollector(string message) : IMessageBusReceiver
{
private int count;
public int Count => Volatile.Read(ref this.count);
public Task ProcessMessage<T>(ComponentBase? sendingComponent, Event triggeredEvent, T? data)
{
if (triggeredEvent is Event.SHOW_ERROR && data is DataErrorMessage error && error.Message == message)
Interlocked.Increment(ref this.count);
return Task.CompletedTask;
}
public Task<MessageBusAnswer<TResult>> ProcessMessageWithResult<TPayload, TResult>(ComponentBase? sendingComponent, Event triggeredEvent, TPayload? data) => Task.FromResult(MessageBusAnswer<TResult>.NONE);
}
}
+69
View File
@@ -0,0 +1,69 @@
using AIStudio.Tools;
using AIStudio.Tools.Services;
using AIStudio.Tools.Services.Indexing;
namespace AIStudio.Tests.Tools;
/// <summary>
/// Checks how a run of the file indexer sorts the files of a folder before it reads the first one.
/// </summary>
/// <remarks>
/// A large folder is indexed in steps of a few minutes, and every step lists the folder anew. What a
/// step counts up front is what its progress starts with: a file counted as to be read lets the bar
/// jump back with every step, and costs the time to read it once more. A file left out wrongly is
/// never indexed at all.
/// </remarks>
[TestFixture]
public sealed class FilePartitionTests
{
private static readonly FileInfo UNCHANGED = new("/tmp/test-data/unchanged.pdf");
private static readonly FileInfo CHANGED = new("/tmp/test-data/changed.pdf");
private static readonly FileInfo ADDED = new("/tmp/test-data/added.pdf");
private static readonly FileInfo SCANNED = new("/tmp/test-data/scanned.pdf");
private static readonly FileInfo SCANNED_AND_RECOGNIZED = new("/tmp/test-data/scanned-and-recognized.pdf");
private static readonly FileInfo FAILED_EARLIER = new("/tmp/test-data/failed-earlier.pdf");
private static readonly FileInfo CHANGED_SINCE_IT_FAILED = new("/tmp/test-data/changed-since-it-failed.pdf");
private static readonly DataSourceEmbeddingFailure EARLIER_FAILURE = new(FAILED_EARLIER.FullName, "The file changed while it was being indexed.", DateTimeOffset.UnixEpoch);
private FilePartition partition = null!;
[SetUp]
public void SortTheFiles()
{
var files = new[] { UNCHANGED, CHANGED, ADDED, SCANNED, SCANNED_AND_RECOGNIZED, FAILED_EARLIER, CHANGED_SINCE_IT_FAILED };
var fingerprints = files.ToDictionary(file => file.FullName, file => $"NOW:{file.Name}", StringComparer.OrdinalIgnoreCase);
var manifest = new DataSourceEmbeddingManifest();
manifest.Files[UNCHANGED.FullName] = Indexed($"NOW:{UNCHANGED.Name}");
manifest.Files[CHANGED.FullName] = Indexed("BEFORE");
manifest.PermanentFailures[SCANNED.FullName] = Unreadable($"NOW:{SCANNED.Name}");
manifest.PermanentFailures[SCANNED_AND_RECOGNIZED.FullName] = Unreadable("BEFORE-OCR");
var pass = new IndexingPass(new DataSourceRunSignature("EMBEDDING-SIGNATURE"));
pass.RecordFailure(FAILED_EARLIER.FullName, $"NOW:{FAILED_EARLIER.Name}", EARLIER_FAILURE);
pass.RecordFailure(CHANGED_SINCE_IT_FAILED.FullName, "BEFORE", EARLIER_FAILURE with { DocumentKey = CHANGED_SINCE_IT_FAILED.FullName });
this.partition = FileSourceIndexer.PartitionFiles(files, fingerprints, manifest, pass);
}
[Test]
public void UnchangedFilesAreCountedButNotRead() =>
Assert.That(this.partition.UnchangedCount, Is.EqualTo(1));
[Test]
public void FilesSkippedForGoodAreNotReadUntilTheyChange() =>
Assert.That(this.partition.StillUnreadable.Select(file => file.Key), Is.EqualTo(new[] { SCANNED.FullName }), "An OCR run over a scanned document changes it, which is exactly when it deserves another attempt.");
[Test]
public void FilesWhichFailedEarlierInThePassAreNotReadUntilTheyChange() =>
Assert.That(this.partition.FailedEarlierInPass, Is.EqualTo(new[] { EARLIER_FAILURE }), "The failure keeps its place in the list, or the list loses what the steps before found.");
[Test]
public void EveryOtherFileIsReadInTheOrderOfTheFolder() =>
Assert.That(this.partition.ToRead.Select(file => file.FullName), Is.EqualTo(new[] { CHANGED.FullName, ADDED.FullName, SCANNED_AND_RECOGNIZED.FullName, CHANGED_SINCE_IT_FAILED.FullName }));
private static EmbeddedFileRecord Indexed(string fingerprint) => new(fingerprint, 1024, DateTimeOffset.UnixEpoch, DateTimeOffset.UnixEpoch, 4);
private static PermanentIndexingFailureRecord Unreadable(string fingerprint) => new(fingerprint, FileExtractionErrorCode.NO_TEXT_EXTRACTED, "No text could be read.", DateTimeOffset.UnixEpoch);
}
+62
View File
@@ -0,0 +1,62 @@
using AIStudio.Tools.Services;
using AIStudio.Tools.Services.Indexing;
namespace AIStudio.Tests.Tools;
/// <summary>
/// Checks when an indexing run carries on the pass of its data source, and when it begins a new one.
/// </summary>
/// <remarks>
/// A pass decides which documents a run leaves alone, because they failed during an earlier step. A
/// run which carries on a pass it should not have left documents unindexed which would work now. A
/// run which begins anew when it should not have reads every failed document once more, with a cloud
/// provider for real money, and tells the user about the same failure once more.
/// </remarks>
[TestFixture]
public sealed class IndexingPassTests
{
private static readonly DataSourceRunSignature SIGNATURE = new("EMBEDDING-SIGNATURE");
[Test]
public void TheNextStepCarriesOnThePass()
{
var pass = new IndexingPass(SIGNATURE);
//
// A request during a step merges into the next step, which then runs for the reason of
// that request:
//
Assert.Multiple(() =>
{
foreach (var refreshMode in new[] { DataSourceEmbeddingRefreshMode.CONTINUATION, DataSourceEmbeddingRefreshMode.WATCHER_HASH_CHECK, DataSourceEmbeddingRefreshMode.INTERVAL_CHECK })
Assert.That(IndexingPass.ContinueOrBegin(pass, refreshMode, SIGNATURE), Is.SameAs(pass), $"RefreshMode={refreshMode}");
});
}
[Test]
public void AManualRetryBeginsANewPass()
{
var pass = new IndexingPass(SIGNATURE);
Assert.That(IndexingPass.ContinueOrBegin(pass, DataSourceEmbeddingRefreshMode.MANUAL_RETRY, SIGNATURE), Is.Not.SameAs(pass), "The user asked for every failed document to be tried again.");
}
[Test]
public void OtherEmbeddingSettingsBeginANewPass()
{
var pass = new IndexingPass(SIGNATURE);
var otherSignature = new DataSourceRunSignature("OTHER-EMBEDDING-SIGNATURE");
var nextPass = IndexingPass.ContinueOrBegin(pass, DataSourceEmbeddingRefreshMode.CONTINUATION, otherSignature);
Assert.Multiple(() =>
{
Assert.That(nextPass, Is.Not.SameAs(pass), "What failed with the old embedding provider may well work with the new one.");
Assert.That(nextPass.RunSignature, Is.EqualTo(otherSignature));
});
}
[Test]
public void ARunWithoutAPassBeginsOne() =>
Assert.That(IndexingPass.ContinueOrBegin(null, DataSourceEmbeddingRefreshMode.CONTINUATION, SIGNATURE).RunSignature, Is.EqualTo(SIGNATURE));
}
+95
View File
@@ -0,0 +1,95 @@
using AIStudio.Settings.DataModel;
using AIStudio.Tools.Services;
namespace AIStudio.Tests.Tools;
/// <summary>
/// Checks what the row of a data source shows while it waits for its next run.
/// </summary>
/// <remarks>
/// A large data source is indexed in steps of a few minutes and waits in the queue between them.
/// Its row must keep the progress and the failures of the steps so far, or the bar jumps back with
/// every step and the failures vanish for minutes. What belongs to a run going on must go, though:
/// a row which names a document while nothing runs, or offers a decision the next run may no longer
/// ask for, tells the user something which is not true.
/// </remarks>
[TestFixture]
public sealed class QueuedStatusTests
{
private static readonly DateTimeOffset NOW = new(2026, 10, 11, 9, 0, 0, TimeSpan.Zero);
private static readonly DataSourceEmbeddingFailure FAILURE = new("/tmp/test-data/report.pdf", "The file could not be read.", NOW);
[Test]
public void AQueuedDataSourceKeepsWhatItsRowShowedOfTheRunsSoFar()
{
var status = RunningStatus();
var queued = DataSourceEmbeddingService.AsQueued(status, DataSource());
Assert.Multiple(() =>
{
Assert.That(queued.State, Is.EqualTo(DataSourceEmbeddingState.QUEUED));
Assert.That(queued.TotalDocuments, Is.EqualTo(status.TotalDocuments));
Assert.That(queued.IndexedDocuments, Is.EqualTo(status.IndexedDocuments));
Assert.That(queued.FailedDocuments, Is.EqualTo(status.FailedDocuments));
Assert.That(queued.PermanentlySkippedDocuments, Is.EqualTo(status.PermanentlySkippedDocuments), "Dropping them made the bar jump back.");
Assert.That(queued.ProgressPercent, Is.EqualTo(status.ProgressPercent), "The bar must not jump back between two steps of the same data source.");
Assert.That(queued.Failures, Is.EqualTo(new[] { FAILURE }));
Assert.That(queued.LastError, Is.EqualTo(status.LastError), "The sentence about the data source as a whole belongs next to its failures.");
Assert.That(queued.LastSyncUtc, Is.EqualTo(NOW));
});
}
[Test]
public void AQueuedDataSourceForgetsWhatBelongsToARunGoingOn()
{
var queued = DataSourceEmbeddingService.AsQueued(RunningStatus(), DataSource());
Assert.Multiple(() =>
{
Assert.That(queued.CurrentDocument, Is.Empty, "Nothing is being read while the data source waits.");
Assert.That(queued.CurrentDocumentBlock, Is.Null);
Assert.That(queued.CurrentDocumentPage, Is.Null);
Assert.That(queued.Attention, Is.EqualTo(DataSourceAttention.NONE), "The next run asks the user anew when the question still applies.");
Assert.That(queued.PendingRemovalCount, Is.Null, "The next run asks the user anew, with the number it finds then.");
Assert.That(queued.VectorStoreUnreadable, Is.False);
});
}
[Test]
public void AQueuedDataSourceShowsItsCurrentName() =>
Assert.That(DataSourceEmbeddingService.AsQueued(RunningStatus(), DataSource()).DataSourceName, Is.EqualTo("Test data"), "The data source may have been renamed since its last run.");
private static DataSourceEmbeddingStatus RunningStatus() => new(
"6f1d6a4e-6a5e-4c62-9a4f-0f2d2c8b7a11",
"Old name",
DataSourceType.LOCAL_DIRECTORY,
DataSourceEmbeddingState.RUNNING,
1200,
400,
3,
"/tmp/test-data/notes.docx",
"Some files could not be read.",
[FAILURE],
PermanentlySkippedDocuments: 7,
CurrentDocumentBlock: 2,
CurrentDocumentPage: 5,
VectorStoreUnreadable: true,
Attention: DataSourceAttention.MASS_REMOVAL_PENDING,
PendingRemovalCount: 42,
LastSyncUtc: NOW);
private static DataSourceLocalDirectory DataSource() => new()
{
Num = 1,
Id = "6f1d6a4e-6a5e-4c62-9a4f-0f2d2c8b7a11",
Name = "Test data",
Description = "Documents used by the tests.",
Type = DataSourceType.LOCAL_DIRECTORY,
EmbeddingId = "b0a4c4d2-1f3e-4f0a-8c9d-5a6b7c8d9e01",
MaxChunkTokenLength = 512,
ChunkOverlapTokenLength = 100,
Path = "/tmp/test-data",
};
}