From 108055954bcd2b9c0d469c292f6f9dd268198a91 Mon Sep 17 00:00:00 2001
From: Ionite <dev@ionite.io>
Date: Tue, 4 Jul 2023 19:31:06 -0400
Subject: [PATCH] Add SharedFolderTests

---
 .../Models/SharedFoldersTests.cs              | 86 +++++++++++++++++++
 .../ReparsePoints/JunctionTests.cs            | 20 +----
 StabilityMatrix.Tests/TempFiles.cs            | 22 +++++
 3 files changed, 109 insertions(+), 19 deletions(-)
 create mode 100644 StabilityMatrix.Tests/Models/SharedFoldersTests.cs
 create mode 100644 StabilityMatrix.Tests/TempFiles.cs

diff --git a/StabilityMatrix.Tests/Models/SharedFoldersTests.cs b/StabilityMatrix.Tests/Models/SharedFoldersTests.cs
new file mode 100644
index 00000000..a372db93
--- /dev/null
+++ b/StabilityMatrix.Tests/Models/SharedFoldersTests.cs
@@ -0,0 +1,86 @@
+using StabilityMatrix.Extensions;
+using StabilityMatrix.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, string>
+        {
+            [SharedFolderType.StableDiffusion] = "models/Stable-diffusion",
+            [SharedFolderType.ESRGAN] = "models/ESRGAN",
+            [SharedFolderType.TextualInversion] = "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}.");
+    }
+}
diff --git a/StabilityMatrix.Tests/ReparsePoints/JunctionTests.cs b/StabilityMatrix.Tests/ReparsePoints/JunctionTests.cs
index 5bb280e3..34615c60 100644
--- a/StabilityMatrix.Tests/ReparsePoints/JunctionTests.cs
+++ b/StabilityMatrix.Tests/ReparsePoints/JunctionTests.cs
@@ -23,25 +23,7 @@ public class JunctionTest
     public void Cleanup()
     {
         if (string.IsNullOrEmpty(tempFolder)) return;
-        DeleteDirectory(tempFolder);
-    }
-
-    // Deletes directory while handling junction folders
-    private static void DeleteDirectory(string directory)
-    {
-        // Enumerate to delete any directory links
-        foreach (var item in Directory.EnumerateDirectories(directory))
-        {
-            var info = new DirectoryInfo(item);
-            if (info.Exists && info.Attributes.HasFlag(FileAttributes.ReparsePoint))
-            {
-                info.Delete();
-            }
-            else
-            {
-                DeleteDirectory(item);
-            }
-        }
+        TempFiles.DeleteDirectory(tempFolder);
     }
 
     [TestMethod]
diff --git a/StabilityMatrix.Tests/TempFiles.cs b/StabilityMatrix.Tests/TempFiles.cs
new file mode 100644
index 00000000..14752b60
--- /dev/null
+++ b/StabilityMatrix.Tests/TempFiles.cs
@@ -0,0 +1,22 @@
+namespace StabilityMatrix.Tests;
+
+public static class TempFiles
+{
+    // Deletes directory while handling junction folders
+    public static void DeleteDirectory(string directory)
+    {
+        // Enumerate to delete any directory links
+        foreach (var item in Directory.EnumerateDirectories(directory))
+        {
+            var info = new DirectoryInfo(item);
+            if (info.Exists && info.Attributes.HasFlag(FileAttributes.ReparsePoint))
+            {
+                info.Delete();
+            }
+            else
+            {
+                DeleteDirectory(item);
+            }
+        }
+    }
+}