using System.IO.Compression; using System.Text; using System.Text.Json; using AIStudio.Chat; using AIStudio.Settings; using AIStudio.Tests.Tools.Databases; using AIStudio.Tools; using AIStudio.Tools.ChatArchive; using AIStudio.Tools.Databases.ProvenanceStore; using AIStudio.Tools.DataBackup; using Microsoft.EntityFrameworkCore; namespace AIStudio.Tests.Tools.DataBackup; /// /// Checks adding the chats of a backup to the chats of another computer. /// /// /// A backup is written on one computer and read on another, whose data folder lies elsewhere. Its /// chats have to arrive with their transcripts, which the chats name by a path relative to the data /// folder, and a chat which exists already must not turn into a second chat with the same ID. /// [TestFixture] [NonParallelizable] public sealed class BackupChatImportTests { private string? previousDataDirectory; private string testDirectory = string.Empty; private string SourceDataDirectory => Path.Combine(this.testDirectory, "source"); private string TargetDataDirectory => Path.Combine(this.testDirectory, "target"); private string ArchivePath => Path.Combine(this.testDirectory, "backup.mwbackup"); [SetUp] public void RememberDataDirectory() { this.previousDataDirectory = SettingsManager.DataDirectory; this.testDirectory = Path.Combine(Path.GetTempPath(), $"ai-studio-backup-chat-import-{Guid.NewGuid():N}"); Directory.CreateDirectory(this.testDirectory); } [TearDown] public void RemoveDataDirectories() { SettingsManager.DataDirectory = this.previousDataDirectory; WorkspaceBehaviour.TestDataDirectory = null; // Whoever reads the cache next must not get the workspaces of this test: WorkspaceBehaviour.InvalidateWorkspaceTreeCache(); try { Directory.Delete(this.testDirectory, true); } catch (IOException) { // A temporary directory we could not remove says nothing about the code under test. } } [Test] public async Task TheChatsOfABackupArriveWithTheirTranscripts() { await using var provenance = await TemporaryProvenanceStore.CreateAsync(); var storedChat = await this.WriteBackupWithAChatAsync(); await UseDataDirectoryAsync(this.TargetDataDirectory); var preview = await ChatArchiveImporter.ReadPreviewAsync(this.ArchivePath, CancellationToken.None); var result = await ChatArchiveImporter.ImportAsync(this.ArchivePath, provenance.Client, ChatArchiveCollisionBehavior.SKIP, null, CancellationToken.None); var chat = await WorkspaceBehaviour.LoadChatAsync(new(storedChat.WorkspaceId, storedChat.ChatId)); var transcript = ((ContentText)chat!.Blocks.Single().Content!).FileAttachments.Single(); var chatsWithEntries = await provenance.Client.GetChatIdsAsync(CancellationToken.None); Assert.Multiple(() => { Assert.That(chatsWithEntries, Is.Empty, "A backup without a provenance store, e.g., of an older version, brings its chats all the same."); Assert.That(preview.Success, Is.True, preview.Issue); Assert.That(preview.Manifest.Workspaces.Select(workspace => workspace.Name), Is.EqualTo(new[] { "Projects" })); Assert.That(preview.Manifest.Workspaces.Single().Chats.Select(listedChat => listedChat.Name), Is.EqualTo(new[] { "Meeting notes" })); Assert.That(result.ImportedChats, Is.EqualTo(1)); Assert.That(result.ImportedWorkspaces, Is.EqualTo(1)); Assert.That(transcript.FilePath, Does.StartWith(WorkspaceBehaviour.GetChatDirectory(storedChat.WorkspaceId, storedChat.ChatId)),"The transcript belongs to the chat on this computer, not to the data folder of the other one."); Assert.That(File.ReadAllText(transcript.FilePath), Does.Contain("Hello and welcome.")); }); } [Test] public async Task AChatWhichExistsAlreadyIsKept() { await using var provenance = await TemporaryProvenanceStore.CreateAsync(); await this.WriteBackupWithAChatAsync(); await UseDataDirectoryAsync(this.TargetDataDirectory); await ChatArchiveImporter.ImportAsync(this.ArchivePath, provenance.Client, ChatArchiveCollisionBehavior.SKIP, null, CancellationToken.None); var secondResult = await ChatArchiveImporter.ImportAsync(this.ArchivePath, provenance.Client, ChatArchiveCollisionBehavior.SKIP, null, CancellationToken.None); Assert.Multiple(() => { Assert.That(secondResult.ImportedChats, Is.Zero); Assert.That(secondResult.SkippedChats, Is.EqualTo(1), "Two chats with one ID would leave the app unable to tell them apart."); }); } [Test] public async Task TheChatsOfABackupArriveWithWhatTheyReceived() { await using var sourceProvenance = await TemporaryProvenanceStore.CreateAsync(); await using var provenance = await TemporaryProvenanceStore.CreateAsync(); var storedChat = await this.WriteBackupWithAChatAsync(sourceProvenance); await UseDataDirectoryAsync(this.TargetDataDirectory); var result = await ChatArchiveImporter.ImportAsync(this.ArchivePath, provenance.Client, ChatArchiveCollisionBehavior.SKIP, null, CancellationToken.None); var entries = await provenance.Client.GetEntriesAsync([storedChat.ProvenanceId!.Value], CancellationToken.None); var chatsWithEntries = await provenance.Client.GetChatIdsAsync(CancellationToken.None); var stagingDirectory = Path.Combine(this.TargetDataDirectory, ".backup-staging"); var stagingLeftovers = Directory.Exists(stagingDirectory) ? Directory.GetFileSystemEntries(stagingDirectory) : Array.Empty(); Assert.Multiple(() => { Assert.That(result.ImportedChats, Is.EqualTo(1)); Assert.That(entries[storedChat.ProvenanceId.Value].Content, Is.EqualTo("Hello and welcome."), "The chat refers to the entry by its id, which the import keeps."); Assert.That(chatsWithEntries, Is.EqualTo(new[] { storedChat.ChatId }), "Another chat's entries of the source computer stay behind."); Assert.That(stagingLeftovers, Is.Empty, "The unpacked store of the backup is removed afterward."); }); } [Test] public async Task AChatAddedOnceMoreTakesItsEntriesAlong() { await using var sourceProvenance = await TemporaryProvenanceStore.CreateAsync(); await using var provenance = await TemporaryProvenanceStore.CreateAsync(); var storedChat = await this.WriteBackupWithAChatAsync(sourceProvenance); await UseDataDirectoryAsync(this.TargetDataDirectory); await ChatArchiveImporter.ImportAsync(this.ArchivePath, provenance.Client, ChatArchiveCollisionBehavior.IMPORT_AS_COPY, null, CancellationToken.None); var secondResult = await ChatArchiveImporter.ImportAsync(this.ArchivePath, provenance.Client, ChatArchiveCollisionBehavior.IMPORT_AS_COPY, null, CancellationToken.None); var chatsWithEntries = await provenance.Client.GetChatIdsAsync(CancellationToken.None); var entryKey = storedChat.ProvenanceId!.Value.ToString("D"); await using var context = provenance.CreateContext(); var copiesOfTheEntry = await context.Entries.CountAsync(entry => entry.Id == entryKey); Assert.Multiple(() => { Assert.That(secondResult.ImportedChats, Is.EqualTo(1)); Assert.That(chatsWithEntries, Has.Count.EqualTo(2), "The copy under its new id gets entries of its own, so deleting one of the two chats leaves the other whole."); Assert.That(chatsWithEntries, Does.Contain(storedChat.ChatId)); Assert.That(copiesOfTheEntry, Is.EqualTo(2), "Both chats refer to the entry by the same id."); }); } [Test] public async Task ABackupOfANewerFormatIsRefused() { var manifest = new DataBackupManifest(DataBackupFormat.FORMAT_VERSION + 1, "99.0.0", "V9", DateTimeOffset.UtcNow, "/data", "/config", false, []); await using (var archiveStream = File.Create(this.ArchivePath)) using (var archive = new ZipArchive(archiveStream, ZipArchiveMode.Create)) { await using var manifestStream = archive.CreateEntry(DataBackupFormat.MANIFEST_ENTRY_NAME).Open(); await JsonSerializer.SerializeAsync(manifestStream, manifest, DataBackupFormat.JSON_OPTIONS); } var preview = await ChatArchiveImporter.ReadPreviewAsync(this.ArchivePath, CancellationToken.None); Assert.Multiple(() => { Assert.That(preview.Success, Is.False, "This version cannot know what a later format holds."); Assert.That(preview.Issue, Is.Not.Empty); }); } [Test] public async Task AZipFileWithoutManifestIsNoBackup() { await using (var archiveStream = File.Create(this.ArchivePath)) using (var archive = new ZipArchive(archiveStream, ZipArchiveMode.Create)) { await using var entryStream = archive.CreateEntry("notes.txt").Open(); await entryStream.WriteAsync(Encoding.UTF8.GetBytes("Not a backup.")); } var preview = await ChatArchiveImporter.ReadPreviewAsync(this.ArchivePath, CancellationToken.None); Assert.That(preview.Success, Is.False); } /// /// Stores a chat with a transcript in a workspace of the source computer, and backs up its data folder. /// /// The provenance store of the source computer, which then holds what the chat received, or null when the backup holds no store. private async Task WriteBackupWithAChatAsync(TemporaryProvenanceStore? sourceProvenance = null) { await UseDataDirectoryAsync(this.SourceDataDirectory); var workspace = await WorkspaceBehaviour.TryCreateWorkspaceAsync("Projects"); Assert.That(workspace.Success, Is.True, "The workspace could not be created."); var workspaceId = workspace.Workspace.WorkspaceId; var chatId = Guid.NewGuid(); var transcriptDirectory = Path.Combine(WorkspaceBehaviour.GetChatDirectory(workspaceId, chatId), "attachments", "transcripts"); Directory.CreateDirectory(transcriptDirectory); var transcriptPath = Path.Combine(transcriptDirectory, "meeting-transcript-0001.md"); await File.WriteAllTextAsync(transcriptPath, "# Transcription: meeting.mp4\n\nHello and welcome."); var chat = new ChatThread { WorkspaceId = workspaceId, ChatId = chatId, Name = "Meeting notes", Blocks = [ new ContentBlock { Time = DateTimeOffset.UtcNow, ContentType = ContentType.TEXT, Content = new ContentText { Text = "Summarize the meeting", FileAttachments = [new ManagedTranscriptAttachment("meeting-transcript-0001.md", transcriptPath, new FileInfo(transcriptPath).Length, "meeting.mp4", false)] }, Role = ChatRole.USER, }, ], }; Assert.That(await WorkspaceBehaviour.StoreChatAsync(chat), Is.True, "The chat could not be stored."); // // What the chat received, next to an entry of a chat the backup does not hold, copied the way // a backup copies the store: // Guid? provenanceId = null; Func>? copyProvenanceStoreAsync = null; if (sourceProvenance is not null) { provenanceId = await sourceProvenance.Client.StoreAsync(chatId, ProvenanceKind.ATTACHMENT, "Hello and welcome.", null, CancellationToken.None); await sourceProvenance.Client.StoreAsync(Guid.NewGuid(), ProvenanceKind.WEB_PAGE, "A page of another chat.", null, CancellationToken.None); var provenanceCopyDirectory = Path.Combine(this.testDirectory, "provenance-copy"); copyProvenanceStoreAsync = async token => { await sourceProvenance.Client.BackupToDirectoryAsync(Path.Combine(provenanceCopyDirectory, "databases", "provenance"), token); return provenanceCopyDirectory; }; } var sources = new DataBackupSources(this.SourceDataDirectory, Path.Combine(this.testDirectory, "config"), "{}", string.Empty, [], copyProvenanceStoreAsync); var manifest = new DataBackupManifest(DataBackupFormat.FORMAT_VERSION, "26.10.1", "V6", DateTimeOffset.UtcNow, this.SourceDataDirectory, string.Empty, false, []); await using (var archiveStream = File.Create(this.ArchivePath)) await DataBackupWriter.WriteAsync(archiveStream, sources, manifest, null, CancellationToken.None); return new(workspaceId, chatId, provenanceId); } /// /// Makes the given folder the data folder, as another computer would have it. /// /// /// The data directory is static state of the whole application, which is why this fixture does /// not run alongside others. The workspace behaviour keeps the roots of the workspaces and of the /// temporary chats fixed, so it gets told about it separately, and its cache starts over. /// private static async Task UseDataDirectoryAsync(string dataDirectory) { Directory.CreateDirectory(dataDirectory); SettingsManager.DataDirectory = dataDirectory; WorkspaceBehaviour.TestDataDirectory = dataDirectory; await WorkspaceBehaviour.ForceReloadWorkspaceTreeAsync(); } /// /// A chat as the source computer stored it. /// /// The workspace of the chat. /// The chat. /// The entry of what the chat received, or null when the backup holds no provenance store. private sealed record StoredChat(Guid WorkspaceId, Guid ChatId, Guid? ProvenanceId = null); }