1 files changed, 31 insertions, 1 deletions
diff --git a/tools/slang-test/test-context.cpp b/tools/slang-test/test-context.cpp
index 06a9847bf..0030d6136 100644
--- a/tools/slang-test/test-context.cpp
+++ b/tools/slang-test/test-context.cpp
@@ -27,6 +27,15 @@ Result TestContext::init()
 
 TestContext::~TestContext()
 {
+    for (auto& pair : m_sharedLibTools)
+    {
+        const auto& tool = pair.Value;
+        if (tool.m_sharedLibrary)
+        {
+            SharedLibrary::unload(tool.m_sharedLibrary);
+        }
+    }
+
     if (m_session)
     {
         spDestroySession(m_session);
@@ -45,7 +54,7 @@ TestContext::InnerMainFunc TestContext::getInnerMainFunc(const String& dirPath,
 
     StringBuilder sharedLibToolBuilder;
     sharedLibToolBuilder.append(name);
-    sharedLibToolBuilder.append("-shared-library");
+    sharedLibToolBuilder.append("-tool");
 
     StringBuilder builder;
     SharedLibrary::appendPlatformFileName(sharedLibToolBuilder.getUnownedSlice(), builder);
@@ -61,3 +70,24 @@ TestContext::InnerMainFunc TestContext::getInnerMainFunc(const String& dirPath,
     m_sharedLibTools.Add(name, tool);
     return tool.m_func;
 }
+
+void TestContext::setInnerMainFunc(const String& name, InnerMainFunc func)
+{
+    SharedLibraryTool* tool = m_sharedLibTools.TryGetValue(name);
+    if (tool)
+    {
+        if (tool->m_sharedLibrary)
+        {
+            SharedLibrary::unload(tool->m_sharedLibrary);
+            tool->m_sharedLibrary = nullptr;
+        }
+
+        tool->m_func = func;
+    }
+    else
+    {
+        SharedLibraryTool tool = {};
+        tool.m_func = func;
+        m_sharedLibTools.Add(name, tool);
+    }
+}