From c741e27ef1f4810df4ffc0a17eb2de23a2873861 Mon Sep 17 00:00:00 2001
From: Kevin Ring <kevin@kotachrome.com>
Date: Tue, 24 Sep 2024 09:39:37 +1000
Subject: [PATCH] Add WaitForTask.

---
 Runtime/WaitForTask.cs       | 25 +++++++++++++++++++++++++
 Runtime/WaitForTask.cs.meta  | 11 +++++++++++
 Tests/TestCesium3DTileset.cs | 20 ++++----------------
 3 files changed, 40 insertions(+), 16 deletions(-)
 create mode 100644 Runtime/WaitForTask.cs
 create mode 100644 Runtime/WaitForTask.cs.meta

diff --git a/Runtime/WaitForTask.cs b/Runtime/WaitForTask.cs
new file mode 100644
index 00000000..94a94efa
--- /dev/null
+++ b/Runtime/WaitForTask.cs
@@ -0,0 +1,25 @@
+using System;
+using UnityEngine;
+
+namespace CesiumForUnity
+{
+    /// <summary>
+    /// A YieldInstruction that can be yielded from a coroutine in order to wait
+    /// until a given task completes.
+    /// </summary>
+    public class WaitForTask : CustomYieldInstruction
+    {
+        private IAsyncResult _task;
+
+        /// <summary>
+        /// Initializes a new instance.
+        /// </summary>
+        /// <param name="task">The task to wait for.</param>
+        public WaitForTask(IAsyncResult task)
+        {
+            this._task = task;
+        }
+
+        public override bool keepWaiting => !this._task.IsCompleted;
+    }
+}
diff --git a/Runtime/WaitForTask.cs.meta b/Runtime/WaitForTask.cs.meta
new file mode 100644
index 00000000..60234ff5
--- /dev/null
+++ b/Runtime/WaitForTask.cs.meta
@@ -0,0 +1,11 @@
+fileFormatVersion: 2
+guid: ca29eff4c876de44b957b0121d454856
+MonoImporter:
+  externalObjects: {}
+  serializedVersion: 2
+  defaultReferences: []
+  executionOrder: 0
+  icon: {instanceID: 0}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 
diff --git a/Tests/TestCesium3DTileset.cs b/Tests/TestCesium3DTileset.cs
index 9cab8e47..bb25b82a 100644
--- a/Tests/TestCesium3DTileset.cs
+++ b/Tests/TestCesium3DTileset.cs
@@ -23,10 +23,7 @@ public IEnumerator SampleHeightMostDetailedWorksWithAnEmptyArrayOfPositions()
 
         Task<CesiumSampleHeightResult> task = tileset.SampleHeightMostDetailed();
 
-        while (!task.IsCompleted)
-        {
-            yield return null;
-        }
+        yield return new WaitForTask(task);
 
         CesiumSampleHeightResult result = task.Result;
         Assert.IsNotNull(result);
@@ -52,10 +49,7 @@ public IEnumerator SampleHeightMostDetailedWorksWithASinglePosition()
 
         Task<CesiumSampleHeightResult> task = tileset.SampleHeightMostDetailed(new double3(-105.1, 40.1, 1.0));
 
-        while (!task.IsCompleted)
-        {
-            yield return null;
-        }
+        yield return new WaitForTask(task);
 
         CesiumSampleHeightResult result = task.Result;
         Assert.IsNotNull(result);
@@ -89,10 +83,7 @@ public IEnumerator SampleHeightMostDetailedWorksWithMultiplePositions()
             new double3(-105.1, 40.1, 1.0),
             new double3(105.1, -40.1, 1.0));
 
-        while (!task.IsCompleted)
-        {
-            yield return null;
-        }
+        yield return new WaitForTask(task);
 
         CesiumSampleHeightResult result = task.Result;
         Assert.IsNotNull(result);
@@ -131,10 +122,7 @@ public IEnumerator SampleHeightMostDetailedIndicatesNotSampledForPositionOutside
         // Somewhere in Sydney, not Melbourne
         Task<CesiumSampleHeightResult> task = tileset.SampleHeightMostDetailed(new double3(151.20972, -33.87100, 1.0));
 
-        while (!task.IsCompleted)
-        {
-            yield return null;
-        }
+        yield return new WaitForTask(task);
 
         CesiumSampleHeightResult result = task.Result;
         Assert.IsNotNull(result);