tg_utils/
hooks.rs

1use schemars::JsonSchema;
2use serde::{Deserialize, Serialize};
3use thiserror::Error;
4
5use cosmwasm_std::{Addr, StdError, StdResult, Storage};
6use cw_storage_plus::Item;
7use tg_bindings::TgradeMsg;
8
9type SubMsg = cosmwasm_std::SubMsg<TgradeMsg>;
10
11// this is copied from cw4
12// TODO: pull into cw_utils as common dep
13#[derive(Serialize, Deserialize, Clone, PartialEq, Eq, JsonSchema, Debug)]
14pub struct HooksResponse {
15    pub hooks: Vec<String>,
16}
17
18#[derive(Error, Debug, PartialEq)]
19pub enum HookError {
20    #[error("{0}")]
21    Std(#[from] StdError),
22
23    #[error("Given address already registered as a hook")]
24    HookAlreadyRegistered {},
25
26    #[error("Given address not registered as a hook")]
27    HookNotRegistered {},
28
29    #[error("You can only unregister yourself from a hook, not other contracts")]
30    OnlyRemoveSelf {},
31}
32
33// store all hook addresses in one item. We cannot have many of them before the contract becomes unusable anyway.
34pub struct Hooks<'a>(Item<'a, Vec<Addr>>);
35
36impl<'a> Hooks<'a> {
37    pub const fn new(hook_key: &'a str) -> Self {
38        Hooks(Item::new(hook_key))
39    }
40
41    pub fn add_hook(&self, storage: &mut dyn Storage, addr: Addr) -> Result<(), HookError> {
42        let mut hooks = self.0.may_load(storage)?.unwrap_or_default();
43        if !hooks.iter().any(|h| h == &addr) {
44            hooks.push(addr);
45        } else {
46            return Err(HookError::HookAlreadyRegistered {});
47        }
48        Ok(self.0.save(storage, &hooks)?)
49    }
50
51    pub fn remove_hook(&self, storage: &mut dyn Storage, addr: Addr) -> Result<(), HookError> {
52        let mut hooks = self.0.load(storage)?;
53        if let Some(p) = hooks.iter().position(|x| x == &addr) {
54            hooks.remove(p);
55        } else {
56            return Err(HookError::HookNotRegistered {});
57        }
58        Ok(self.0.save(storage, &hooks)?)
59    }
60
61    pub fn list_hooks(&self, storage: &dyn Storage) -> StdResult<Vec<String>> {
62        let hooks = self.0.may_load(storage)?.unwrap_or_default();
63        Ok(hooks.into_iter().map(String::from).collect())
64    }
65
66    pub fn prepare_hooks<F: Fn(Addr) -> StdResult<SubMsg>>(
67        &self,
68        storage: &dyn Storage,
69        prep: F,
70    ) -> StdResult<Vec<SubMsg>> {
71        self.0
72            .may_load(storage)?
73            .unwrap_or_default()
74            .into_iter()
75            .map(prep)
76            .collect()
77    }
78}
79
80#[cfg(test)]
81mod test {
82    use super::*;
83    use cosmwasm_std::testing::mock_dependencies;
84    use cosmwasm_std::{coins, BankMsg, CosmosMsg, Deps};
85
86    const HOOKS: Hooks = Hooks::new("hooks");
87
88    fn assert_count(deps: Deps, expected: usize) {
89        let hooks = HOOKS.list_hooks(deps.storage).unwrap();
90        assert_eq!(hooks.len(), expected);
91    }
92
93    #[test]
94    fn add_and_remove_hooks() {
95        let mut deps = mock_dependencies();
96        assert_count(deps.as_ref(), 0);
97
98        // add a new hook
99        let first = Addr::unchecked("first");
100        HOOKS
101            .add_hook(deps.as_mut().storage, first.clone())
102            .unwrap();
103        assert_count(deps.as_ref(), 1);
104
105        // cannot add twice
106        let err = HOOKS
107            .add_hook(deps.as_mut().storage, first.clone())
108            .unwrap_err();
109        assert_eq!(err, HookError::HookAlreadyRegistered {});
110        assert_count(deps.as_ref(), 1);
111
112        // add a different hook
113        let bar = Addr::unchecked("bar");
114        HOOKS.add_hook(deps.as_mut().storage, bar).unwrap();
115        assert_count(deps.as_ref(), 2);
116
117        // cannot remove a non-registered hook
118        let boom = Addr::unchecked("boom");
119        let err = HOOKS.remove_hook(deps.as_mut().storage, boom).unwrap_err();
120        assert_eq!(err, HookError::HookNotRegistered {});
121        assert_count(deps.as_ref(), 2);
122
123        // can remove one of the existing hooks
124        HOOKS.remove_hook(deps.as_mut().storage, first).unwrap();
125        assert_count(deps.as_ref(), 1);
126    }
127
128    #[test]
129    fn prepare_hook() {
130        let payout = |addr: Addr| {
131            Ok(SubMsg::new(BankMsg::Send {
132                to_address: addr.into(),
133                amount: coins(12345, "bonus"),
134            }))
135        };
136        let mut deps = mock_dependencies();
137        let storage = deps.as_mut().storage;
138
139        HOOKS.add_hook(storage, Addr::unchecked("some")).unwrap();
140        HOOKS.add_hook(storage, Addr::unchecked("one")).unwrap();
141
142        let mut msgs = HOOKS.prepare_hooks(storage, payout).unwrap();
143        assert_eq!(msgs.len(), 2);
144        // get the last message
145        match msgs.pop().unwrap().msg {
146            CosmosMsg::Bank(BankMsg::Send { to_address, amount }) => {
147                assert_eq!(to_address.as_str(), "one");
148                assert_eq!(amount, coins(12345, "bonus"));
149            }
150            _ => panic!("bad message"),
151        }
152    }
153}
tg_utils/hooks.rs

tg_utils/
hooks.rs
