From 8996fbebd7fb2f06dbec28d7014542134b090280 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Wo=C5=BAniak?= <jnwznk@gmail.com>
Date: Mon, 5 Feb 2024 15:51:56 +0100
Subject: [PATCH] test: Add sudo to custom example

---
 examples/contracts/custom/src/contract.rs     | 31 ++++++++++++++++---
 .../contracts/custom/src/multitest/tests.rs   |  5 +++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/examples/contracts/custom/src/contract.rs b/examples/contracts/custom/src/contract.rs
index df0806a7..4890aa2d 100644
--- a/examples/contracts/custom/src/contract.rs
+++ b/examples/contracts/custom/src/contract.rs
@@ -1,5 +1,6 @@
 use cosmwasm_std::{CosmosMsg, QueryRequest, Response, StdResult};
-use sylvia::types::{ExecCtx, InstantiateCtx, QueryCtx};
+use cw_storage_plus::Item;
+use sylvia::types::{ExecCtx, InstantiateCtx, QueryCtx, SudoCtx};
 use sylvia::{contract, schemars};
 
 #[cfg(not(feature = "library"))]
@@ -7,7 +8,9 @@ use sylvia::entry_points;
 
 use crate::messages::{CountResponse, CounterMsg, CounterQuery};
 
-pub struct CustomContract;
+pub struct CustomContract {
+    pub(crate) sudo_counter: Item<'static, u64>,
+}
 
 #[cfg_attr(not(feature = "library"), entry_points)]
 #[contract]
@@ -15,14 +18,17 @@ pub struct CustomContract;
 #[sv::custom(query=CounterQuery, msg=CounterMsg)]
 impl CustomContract {
     pub const fn new() -> Self {
-        Self
+        Self {
+            sudo_counter: Item::new("sudo_counter"),
+        }
     }
 
     #[msg(instantiate)]
     pub fn instantiate(
         &self,
-        _ctx: InstantiateCtx<CounterQuery>,
+        ctx: InstantiateCtx<CounterQuery>,
     ) -> StdResult<Response<CounterMsg>> {
+        self.sudo_counter.save(ctx.deps.storage, &0)?;
         Ok(Response::default())
     }
 
@@ -42,4 +48,21 @@ impl CustomContract {
 
         Ok(resp)
     }
+
+    #[msg(query)]
+    pub fn sudo_counter(&self, ctx: QueryCtx<CounterQuery>) -> StdResult<CountResponse> {
+        let count = self.sudo_counter.load(ctx.deps.storage)?;
+
+        Ok(CountResponse { count })
+    }
+
+    #[msg(sudo)]
+    pub fn increment_sudo_counter(
+        &self,
+        ctx: SudoCtx<CounterQuery>,
+    ) -> StdResult<Response<CounterMsg>> {
+        self.sudo_counter
+            .update(ctx.deps.storage, |value| -> StdResult<_> { Ok(value + 1) })?;
+        Ok(Response::new())
+    }
 }
diff --git a/examples/contracts/custom/src/multitest/tests.rs b/examples/contracts/custom/src/multitest/tests.rs
index 566ee7ff..6edc9eb9 100644
--- a/examples/contracts/custom/src/multitest/tests.rs
+++ b/examples/contracts/custom/src/multitest/tests.rs
@@ -24,4 +24,9 @@ fn test_custom() {
 
     let count = contract.query_custom().unwrap().count;
     assert_eq!(count, 1);
+
+    contract.increment_sudo_counter().unwrap();
+
+    let count = contract.sudo_counter().unwrap().count;
+    assert_eq!(count, 1);
 }