Ionite
1 year ago
3 changed files with 109 additions and 19 deletions
@ -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}."); |
||||||
|
} |
||||||
|
} |
@ -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); |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
} |
Loading…
Reference in new issue