Add show Gizmos argument to TakeScreenshotAttribute

nowsprinting · Oct 30, 2023 · b6dd4fd · b6dd4fd
1 parent 5bc1f8c
commit b6dd4fd
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 3 deletions.
diff --git a/Runtime/Attributes/TakeScreenshotAttribute.cs b/Runtime/Attributes/TakeScreenshotAttribute.cs
@@ -5,9 +5,10 @@
 using System.Collections;
 using NUnit.Framework;
 using NUnit.Framework.Interfaces;
-using TestHelper.Utils;
+using TestHelper.RuntimeInternals;
 using UnityEngine;
 using UnityEngine.TestTools;
+using ScreenshotHelper = TestHelper.Utils.ScreenshotHelper;
 
 namespace TestHelper.Attributes
 {
@@ -21,6 +22,7 @@ public class TakeScreenshotAttribute : NUnitAttribute, IOuterUnityTestAction
         private readonly string _filename;
         private readonly int _superSize;
         private readonly ScreenCapture.StereoScreenCaptureMode _stereoCaptureMode;
+        private readonly bool _gizmos;
 
         /// <summary>
         /// Take a screenshot and save it to file after running the test.
@@ -38,17 +40,19 @@ public class TakeScreenshotAttribute : NUnitAttribute, IOuterUnityTestAction
         /// <param name="filename">Filename to store screenshot.</param>
         /// <param name="superSize">The factor to increase resolution with.</param>
         /// <param name="stereoCaptureMode">The eye texture to capture when stereo rendering is enabled.</param>
+        /// <param name="gizmos">True: show Gizmos on GameView</param>
         public TakeScreenshotAttribute(
             string directory = null,
             string filename = null,
             int superSize = 1,
-            ScreenCapture.StereoScreenCaptureMode stereoCaptureMode = ScreenCapture.StereoScreenCaptureMode.LeftEye
-        )
+            ScreenCapture.StereoScreenCaptureMode stereoCaptureMode = ScreenCapture.StereoScreenCaptureMode.LeftEye,
+            bool gizmos = false)
         {
             _directory = directory;
             _filename = filename;
             _superSize = superSize;
             _stereoCaptureMode = stereoCaptureMode;
+            _gizmos = gizmos;
         }
 
         /// <inheritdoc />
@@ -60,7 +64,19 @@ public IEnumerator BeforeTest(ITest test)
         /// <inheritdoc />
         public IEnumerator AfterTest(ITest test)
         {
+            var beforeGizmos = false;
+            if (_gizmos)
+            {
+                beforeGizmos = GameViewControlHelper.GetGizmos();
+                GameViewControlHelper.SetGizmos(true);
+            }
+
             yield return ScreenshotHelper.TakeScreenshot(_directory, _filename, _superSize, _stereoCaptureMode);
+
+            if (_gizmos)
+            {
+                GameViewControlHelper.SetGizmos(beforeGizmos);
+            }
         }
     }
 }
diff --git a/Tests/Runtime/Attributes/TakeScreenshotAttributeTest.cs b/Tests/Runtime/Attributes/TakeScreenshotAttributeTest.cs
@@ -141,5 +141,25 @@ public void AttachWithFilename_TakeScreenshotAndSaveToSpecifyPath_AfterRunningTe
                 $"{nameof(AttachWithFilename_TakeScreenshotAndSaveToSpecifyPath)}.png");
             Assert.That(path, Does.Exist);
         }
+
+        private class GizmoDemo : MonoBehaviour
+        {
+            private void OnDrawGizmos()
+            {
+                Gizmos.color = Color.red;
+                Gizmos.DrawSphere(transform.position, 0.2f);
+            }
+        }
+
+        [Test]
+        [LoadScene(TestScene)]
+        [TakeScreenshot(gizmos: true)]
+        public void AttachWithGizmos_TakeScreenshotWithGizmos()
+        {
+            var gameObject = GameObject.CreatePrimitive(PrimitiveType.Cube);
+            gameObject.AddComponent<GizmoDemo>();
+
+            // Take screenshot after running the test.
+        }
     }
 }