From 1ea6b78b892a43d96801fee39a09aa8ab6a0ea62 Mon Sep 17 00:00:00 2001
From: nan01ab <yjcc201374@outlook.com>
Date: Sat, 25 Jan 2025 00:28:04 +0800
Subject: [PATCH] Feature: add unit tests for NeedSnapshot and
 NeedApplicationEngine

---
 .../Native/UT_ContractMethodAttribute.cs      | 46 +++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/tests/Neo.UnitTests/SmartContract/Native/UT_ContractMethodAttribute.cs b/tests/Neo.UnitTests/SmartContract/Native/UT_ContractMethodAttribute.cs
index 780c270b6b..abf7cc1fe5 100644
--- a/tests/Neo.UnitTests/SmartContract/Native/UT_ContractMethodAttribute.cs
+++ b/tests/Neo.UnitTests/SmartContract/Native/UT_ContractMethodAttribute.cs
@@ -10,7 +10,10 @@
 // modifications are permitted.
 
 using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Neo.Persistence;
+using Neo.SmartContract;
 using Neo.SmartContract.Native;
+using System.Reflection;
 
 namespace Neo.UnitTests.SmartContract.Native
 {
@@ -28,5 +31,48 @@ public void TestConstructorOneArg()
 
             Assert.AreEqual(Hardfork.HF_Aspidochelone, arg.ActiveIn);
         }
+
+        class NeedSnapshot
+        {
+            [ContractMethod]
+            public bool MethodReadOnlyStoreView(IReadOnlyStoreView view) => view is null;
+
+            [ContractMethod]
+            public bool MethodDataCache(DataCache dataCache) => dataCache is null;
+        }
+
+        class NoNeedSnapshot
+        {
+            [ContractMethod]
+            public bool MethodTwo(ApplicationEngine engine, UInt160 account)
+                => engine is null || account is null;
+
+            [ContractMethod]
+            public bool MethodOne(ApplicationEngine engine) => engine is null;
+        }
+
+        [TestMethod]
+        public void TestNeedSnapshot()
+        {
+            var flags = BindingFlags.Instance | BindingFlags.Static | BindingFlags.NonPublic | BindingFlags.Public;
+            foreach (var member in typeof(NeedSnapshot).GetMembers(flags))
+            {
+                foreach (var attribute in member.GetCustomAttributes<ContractMethodAttribute>())
+                {
+                    var metadata = new ContractMethodMetadata(member, attribute);
+                    Assert.IsTrue(metadata.NeedSnapshot);
+                }
+            }
+
+            foreach (var member in typeof(NoNeedSnapshot).GetMembers(flags))
+            {
+                foreach (var attribute in member.GetCustomAttributes<ContractMethodAttribute>())
+                {
+                    var metadata = new ContractMethodMetadata(member, attribute);
+                    Assert.IsFalse(metadata.NeedSnapshot);
+                    Assert.IsTrue(metadata.NeedApplicationEngine);
+                }
+            }
+        }
     }
 }