AI-Studio/app/Tests/Tools/DataBackup/BackupChatImportTests.cs

using System.IO.Compression;
using System.Text;
using System.Text.Json;

using AIStudio.Chat;
using AIStudio.Settings;
using AIStudio.Tools;
using AIStudio.Tools.ChatArchive;
using AIStudio.Tools.DataBackup;

namespace AIStudio.Tests.Tools.DataBackup;

/// <summary>
/// Checks adding the chats of a backup to the chats of another computer.
/// </summary>
/// <remarks>
/// 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.
/// </remarks>
[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()
    {
        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, ChatArchiveCollisionBehavior.SKIP, null, CancellationToken.None);
        var chat = await WorkspaceBehaviour.LoadChatAsync(new(storedChat.WorkspaceId, storedChat.ChatId));
        var transcript = ((ContentText)chat!.Blocks.Single().Content!).FileAttachments.Single();

        Assert.Multiple(() =>
        {
            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 this.WriteBackupWithAChatAsync();
        await UseDataDirectoryAsync(this.TargetDataDirectory);

        await ChatArchiveImporter.ImportAsync(this.ArchivePath, ChatArchiveCollisionBehavior.SKIP, null, CancellationToken.None);
        var secondResult = await ChatArchiveImporter.ImportAsync(this.ArchivePath, 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 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);
    }

    /// <summary>
    /// Stores a chat with a transcript in a workspace of the source computer, and backs up its data folder.
    /// </summary>
    private async Task<StoredChat> WriteBackupWithAChatAsync()
    {
        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.");

        var sources = new DataBackupSources(this.SourceDataDirectory, Path.Combine(this.testDirectory, "config"), "{}", string.Empty, []);
        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);
    }

    /// <summary>
    /// Makes the given folder the data folder, as another computer would have it.
    /// </summary>
    /// <remarks>
    /// 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.
    /// </remarks>
    private static async Task UseDataDirectoryAsync(string dataDirectory)
    {
        Directory.CreateDirectory(dataDirectory);
        SettingsManager.DataDirectory = dataDirectory;
        WorkspaceBehaviour.TestDataDirectory = dataDirectory;
        await WorkspaceBehaviour.ForceReloadWorkspaceTreeAsync();
    }

    /// <summary>
    /// A chat as the source computer stored it.
    /// </summary>
    private sealed record StoredChat(Guid WorkspaceId, Guid ChatId);
}