working implementation with async

undreamai · Feb 26, 2024 · d99f7e4 · d99f7e4
1 parent e10b86f
commit d99f7e4
Showing 1 changed file with 22 additions and 5 deletions.
diff --git a/Tests/Runtime/TestLLM.cs b/Tests/Runtime/TestLLM.cs
@@ -3,6 +3,9 @@
 using UnityEngine;
 using System.Threading.Tasks;
 using System.Collections.Generic;
+using System;
+using System.Collections;
+using UnityEngine.TestTools;
 
 namespace LLMUnityTests
 {
@@ -39,7 +42,7 @@ public class TestLLM
         LLMNoAwake llm;
         int port = 15555;
         string AIReply = ":::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::";
-
+        Exception error = null;
 
         public TestLLM()
         {
@@ -67,13 +70,14 @@ public async Task Init()
             llm.seed = 0;
             llm.stream = false;
             llm.numPredict = 128;
+            llm.parallelPrompts = 1;
 
             gameObject.SetActive(true);
         }
 
-        [Test]
-        public async void RunTests()
+        public async Task RunTests()
         {
+            error = null;
             try
             {
                 Assert.That(!llm.IsPortInUse());
@@ -88,9 +92,22 @@ public async void RunTests()
                 await llm.SetPrompt("You are");
                 TestInitParameters();
             }
-            finally
+            catch  (Exception e)
+            {
+                error = e;
+            }
+        }
+
+        [UnityTest]
+        public IEnumerator RunTestsWait()
+        {
+            Task task = RunTests();
+            while (!task.IsCompleted) yield return null;
+            llm.CallOnDestroy();
+            if (error != null)
             {
-                llm.CallOnDestroy();
+                Debug.LogError(error.ToString());
+                throw (error);
             }
         }