1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163

use crate::CwEnv;
use crate::{
    error::BootError, index_response::IndexResponse, state::StateInterface, tx_handler::TxResponse,
};
use cosmwasm_std::{Addr, Coin, CustomQuery, Empty};
use cw_multi_test::Contract as TestContract;
use schemars::JsonSchema;
use serde::{de::DeserializeOwned, Serialize};
use std::fmt::Debug;

/// An instance of a contract. Contains references to the execution environment (chain) and a local state (state)
/// The state is used to store contract addresses/code-ids
#[derive(Clone)]
pub struct Contract<Chain: CwEnv> {
    /// ID of the contract, used to retrieve addr/code-id
    pub id: String,
    pub(crate) source: ContractCodeReference,
    /// chain object that handles tx execution and queries.
    pub(crate) chain: Chain,
}

#[derive(Default)]
pub struct ContractCodeReference<ExecT = Empty, QueryT = Empty>
where
    ExecT: Clone + Debug + PartialEq + JsonSchema + DeserializeOwned + 'static,
    QueryT: CustomQuery + DeserializeOwned + 'static,
{
    pub wasm_code_path: Option<String>,
    pub contract_endpoints: Option<Box<dyn TestContract<ExecT, QueryT>>>,
}

impl Clone for ContractCodeReference {
    fn clone(&self) -> Self {
        Self {
            wasm_code_path: self.wasm_code_path.clone(),
            contract_endpoints: None,
        }
    }
}

/// Expose chain and state function to call them on the contract
impl<Chain: CwEnv + Clone> Contract<Chain> {
    pub fn new(id: impl ToString, chain: Chain) -> Self {
        Contract {
            id: id.to_string(),
            chain,
            source: ContractCodeReference::default(),
        }
    }

    /// `get_chain` instead of `chain` to disambiguate from the std prelude .chain() method.
    pub fn get_chain(&self) -> &Chain {
        &self.chain
    }

    pub fn with_wasm_path(mut self, path: impl ToString) -> Self {
        self.source.wasm_code_path = Some(path.to_string());
        self
    }

    pub fn with_mock(mut self, mock_contract: Box<dyn TestContract<Empty, Empty>>) -> Self {
        self.source.contract_endpoints = Some(mock_contract);
        self
    }

    pub fn set_mock(&mut self, mock_contract: Box<dyn TestContract<Empty, Empty>>) {
        self.source.contract_endpoints = Some(mock_contract);
    }

    /// Sets the address of the contract in the local state
    pub fn with_address(self, address: Option<&Addr>) -> Self {
        if let Some(address) = address {
            self.set_address(address)
        }
        self
    }

    // Chain interfaces
    pub fn execute<E: Serialize + Debug>(
        &self,
        msg: &E,
        coins: Option<&[Coin]>,
    ) -> Result<TxResponse<Chain>, BootError> {
        log::info!("Executing {:#?} on {}", msg, self.id);
        let resp = self
            .chain
            .execute(msg, coins.unwrap_or(&[]), &self.address()?);
        log::debug!("execute response: {:?}", resp);
        resp.map_err(Into::into)
    }

    pub fn instantiate<I: Serialize + Debug>(
        &self,
        msg: &I,
        admin: Option<&Addr>,
        coins: Option<&[Coin]>,
    ) -> Result<TxResponse<Chain>, BootError> {
        log::info!("Instantiating {} with msg {:#?}", self.id, msg);
        let resp = self
            .chain
            .instantiate(
                self.code_id()?,
                msg,
                Some(&self.id),
                admin,
                coins.unwrap_or(&[]),
            )
            .map_err(Into::into)?;
        let contract_address = resp.instantiated_contract_address()?;
        self.set_address(&contract_address);
        log::info!("Instantiated {} with address {}", self.id, contract_address);
        log::debug!("Instantiate response: {:?}", resp);
        Ok(resp)
    }

    pub fn upload(&mut self) -> Result<TxResponse<Chain>, BootError> {
        log::info!("Uploading {}", self.id);
        let resp = self.chain.upload(&mut self.source).map_err(Into::into)?;
        let code_id = resp.uploaded_code_id()?;
        self.set_code_id(code_id);
        log::info!("uploaded {} with code id {}", self.id, code_id);
        log::debug!("Upload response: {:?}", resp);
        Ok(resp)
    }

    pub fn query<Q: Serialize + Debug, T: Serialize + DeserializeOwned + Debug>(
        &self,
        query_msg: &Q,
    ) -> Result<T, BootError> {
        log::info!("Querying {:#?} on {}", query_msg, self.id);
        let resp = self
            .chain
            .query(query_msg, &self.address()?)
            .map_err(Into::into)?;
        log::debug!("Query response: {:?}", resp);
        Ok(resp)
    }

    pub fn migrate<M: Serialize + Debug>(
        &self,
        migrate_msg: &M,
        new_code_id: u64,
    ) -> Result<TxResponse<Chain>, BootError> {
        log::info!("Migrating {:?} to code_id {}", self.id, new_code_id);
        self.chain
            .migrate(migrate_msg, new_code_id, &self.address()?)
            .map_err(Into::into)
    }

    // State interfaces
    pub fn address(&self) -> Result<Addr, BootError> {
        self.chain.state().get_address(&self.id)
    }
    pub fn code_id(&self) -> Result<u64, BootError> {
        self.chain.state().get_code_id(&self.id)
    }
    pub fn set_address(&self, address: &Addr) {
        self.chain.state().set_address(&self.id, address)
    }
    pub fn set_code_id(&self, code_id: u64) {
        self.chain.state().set_code_id(&self.id, code_id)
    }
}