StabilityMatrix/StabilityMatrix.Tests/Models/SharedFoldersTests.cs

using StabilityMatrix.Core.Extensions;
using StabilityMatrix.Core.Helper;
using StabilityMatrix.Core.Models;

namespace StabilityMatrix.Tests.Models;

[TestClass]
public class SharedFoldersTests
{
    private string tempFolder = string.Empty;
    private string TempModelsFolder => Path.Combine(tempFolder, "models");
    private string TempPackageFolder => Path.Combine(tempFolder, "package");
    
    private readonly Dictionary<SharedFolderType, string> sampleDefinitions = new()
    {
        [SharedFolderType.StableDiffusion] = "models/Stable-diffusion",
        [SharedFolderType.ESRGAN] = "models/ESRGAN",
        [SharedFolderType.TextualInversion] = "embeddings",
    };

    [TestInitialize]
    public void Initialize()
    {
        tempFolder = Path.GetTempFileName();
        File.Delete(tempFolder);
        Directory.CreateDirectory(tempFolder);
    }

    [TestCleanup]
    public void Cleanup()
    {
        if (string.IsNullOrEmpty(tempFolder)) return;
        TempFiles.DeleteDirectory(tempFolder);
    }

    private void CreateSampleJunctions()
    {
        var definitions = new Dictionary<SharedFolderType, IReadOnlyList<string>>
        {
            [SharedFolderType.StableDiffusion] = new[] {"models/Stable-diffusion"},
            [SharedFolderType.ESRGAN] = new[] {"models/ESRGAN"},
            [SharedFolderType.TextualInversion] = new[] {"embeddings"},
        };
        SharedFolders.SetupLinks(definitions, TempModelsFolder, TempPackageFolder);
    }

    [TestMethod]
    public void SetupLinks_CreatesJunctions()
    {
        CreateSampleJunctions();
        
        // Check model folders
        foreach (var (folderType, relativePath) in sampleDefinitions)
        {
            var packagePath = Path.Combine(TempPackageFolder, relativePath);
            var modelFolder = Path.Combine(TempModelsFolder, folderType.GetStringValue());
            // Should exist and be a junction
            Assert.IsTrue(Directory.Exists(packagePath), $"Package folder {packagePath} does not exist.");
            var info = new DirectoryInfo(packagePath);
            Assert.IsTrue(info.Attributes.HasFlag(FileAttributes.ReparsePoint), $"Package folder {packagePath} is not a junction.");
            // Check junction target should be in models folder
            Assert.AreEqual(modelFolder, info.LinkTarget, $"Package folder {packagePath} does not point to {modelFolder}.");
        }
    }

    [TestMethod]
    public void SetupLinks_CanDeleteJunctions()
    {
        CreateSampleJunctions();
        
        var modelFolder = Path.Combine(tempFolder, "models", SharedFolderType.StableDiffusion.GetStringValue());
        var packagePath = Path.Combine(tempFolder, "package", sampleDefinitions[SharedFolderType.StableDiffusion]);
        
        // Write a file to a model folder
        File.Create(Path.Combine(modelFolder, "AFile")).Close();
        Assert.IsTrue(File.Exists(Path.Combine(modelFolder, "AFile")), $"File should exist in {modelFolder}.");
        // Should exist in the package folder
        Assert.IsTrue(File.Exists(Path.Combine(packagePath, "AFile")), $"File should exist in {packagePath}.");
        
        // Now delete the junction
        Directory.Delete(packagePath, false);
        Assert.IsFalse(Directory.Exists(packagePath), $"Package folder {packagePath} should not exist.");
        
        // The file should still exist in the model folder
        Assert.IsTrue(File.Exists(Path.Combine(modelFolder, "AFile")), $"File should exist in {modelFolder}.");
    }
}