Repository: DA0-DA0/dao-contracts
Pull request #278 (2022-04-28 23:20:56)

Vt rewards distributor


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index e569e30d..e4960c0f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -990,6 +990,25 @@ dependencies = [
  "thiserror",
 ]
 
+[[package]]
+name = "stake-cw20-reward-distributor"
+version = "0.1.0"
+dependencies = [
+ "cosmwasm-schema",
+ "cosmwasm-std",
+ "cosmwasm-storage",
+ "cw-multi-test",
+ "cw-storage-plus",
+ "cw-utils",
+ "cw2",
+ "cw20",
+ "cw20-base",
+ "schemars",
+ "serde",
+ "stake-cw20",
+ "thiserror",
+]
+
 [[package]]
 name = "static_assertions"
 version = "1.1.0"
diff --git a/contracts/stake-cw20-reward-distributor/.cargo/config b/contracts/stake-cw20-reward-distributor/.cargo/config
new file mode 100644
index 00000000..336b618a
--- /dev/null
+++ b/contracts/stake-cw20-reward-distributor/.cargo/config
@@ -0,0 +1,4 @@
+[alias]
+wasm = "build --release --target wasm32-unknown-unknown"
+unit-test = "test --lib"
+schema = "run --example schema"
diff --git a/contracts/stake-cw20-reward-distributor/Cargo.toml b/contracts/stake-cw20-reward-distributor/Cargo.toml
new file mode 100644
index 00000000..4bbaa53d
--- /dev/null
+++ b/contracts/stake-cw20-reward-distributor/Cargo.toml
@@ -0,0 +1,46 @@
+[package]
+name = "stake-cw20-reward-distributor"
+version = "0.1.0"
+edition = "2018"
+authors = ["Vernon Johnson <vtj2105@columbia.edu>"]
+
+exclude = [
+  # Those files are rust-optimizer artifacts. You might want to commit them for convenience but they should not be part of the source code publication.
+  "contract.wasm",
+  "hash.txt",
+]
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[lib]
+crate-type = ["cdylib", "rlib"]
+
+[features]
+# for more explicit tests, cargo test --features=backtraces
+backtraces = ["cosmwasm-std/backtraces"]
+# use library feature to disable all instantiate/execute/query exports
+library = []
+
+[package.metadata.scripts]
+optimize = """docker run --rm -v "$(pwd)":/code \
+  --mount type=volume,source="$(basename "$(pwd)")_cache",target=/code/target \
+  --mount type=volume,source=registry_cache,target=/usr/local/cargo/registry \
+  cosmwasm/rust-optimizer:0.12.5
+"""
+
+[dependencies]
+cosmwasm-std = { version = "1.0.0-beta" }
+cosmwasm-storage = { version = "1.0.0-beta" }
+cw-storage-plus = "0.13"
+cw2 = "0.13"
+cw20 = { version = "0.13" }
+cw-utils = "0.13"
+schemars = "0.8"
+cw20-base = {  version = "0.13", features = ["library"] }
+stake-cw20 = { path = "../stake-cw20", features = ["library"]}
+serde = { version = "1.0", default-features = false, features = ["derive"] }
+thiserror = { version = "1.0" }
+
+[dev-dependencies]
+cosmwasm-schema = { version = "1.0.0-beta" }
+cw-multi-test = "0.13"
diff --git a/contracts/stake-cw20-reward-distributor/README.md b/contracts/stake-cw20-reward-distributor/README.md
new file mode 100644
index 00000000..3da02559
--- /dev/null
+++ b/contracts/stake-cw20-reward-distributor/README.md
@@ -0,0 +1,3 @@
+# stake-cw20-reward-distributor
+
+A contract to fund staking contracts
diff --git a/contracts/stake-cw20-reward-distributor/examples/schema.rs b/contracts/stake-cw20-reward-distributor/examples/schema.rs
new file mode 100644
index 00000000..9fd459a9
--- /dev/null
+++ b/contracts/stake-cw20-reward-distributor/examples/schema.rs
@@ -0,0 +1,16 @@
+use std::env::current_dir;
+use std::fs::create_dir_all;
+
+use cosmwasm_schema::{export_schema, remove_schemas, schema_for};
+use stake_cw20_reward_distributor::msg::{ExecuteMsg, InstantiateMsg, QueryMsg};
+
+fn main() {
+    let mut out_dir = current_dir().unwrap();
+    out_dir.push("schema");
+    create_dir_all(&out_dir).unwrap();
+    remove_schemas(&out_dir).unwrap();
+
+    export_schema(&schema_for!(InstantiateMsg), &out_dir);
+    export_schema(&schema_for!(ExecuteMsg), &out_dir);
+    export_schema(&schema_for!(QueryMsg), &out_dir);
+}
diff --git a/contracts/stake-cw20-reward-distributor/schema/execute_msg.json b/contracts/stake-cw20-reward-distributor/schema/execute_msg.json
new file mode 100644
index 00000000..021e5d10
--- /dev/null
+++ b/contracts/stake-cw20-reward-distributor/schema/execute_msg.json
@@ -0,0 +1,56 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "title": "ExecuteMsg",
+  "oneOf": [
+    {
+      "type": "object",
+      "required": [
+        "update_config"
+      ],
+      "properties": {
+        "update_config": {
+          "type": "object",
+          "required": [
+            "owner",
+            "reward_rate",
+            "reward_token",
+            "staking_addr"
+          ],
+          "properties": {
+            "owner": {
+              "type": "string"
+            },
+            "reward_rate": {
+              "$ref": "#/definitions/Uint128"
+            },
+            "reward_token": {
+              "type": "string"
+            },
+            "staking_addr": {
+              "type": "string"
+            }
+          }
+        }
+      },
+      "additionalProperties": false
+    },
+    {
+      "type": "object",
+      "required": [
+        "distribute"
+      ],
+      "properties": {
+        "distribute": {
+          "type": "object"
+        }
+      },
+      "additionalProperties": false
+    }
+  ],
+  "definitions": {
+    "Uint128": {
+      "description": "A thin wrapper around u128 that is using strings for JSON encoding/decoding, such that the full u128 range can be used for clients that convert JSON numbers to floats, like JavaScript and jq.\n\n# Examples\n\nUse `from` to create instances of this and `u128` to get the value out:\n\n``` # use cosmwasm_std::Uint128; let a = Uint128::from(123u128); assert_eq!(a.u128(), 123);\n\nlet b = Uint128::from(42u64); assert_eq!(b.u128(), 42);\n\nlet c = Uint128::from(70u32); assert_eq!(c.u128(), 70); ```",
+      "type": "string"
+    }
+  }
+}
diff --git a/contracts/stake-cw20-reward-distributor/schema/instantiate_msg.json b/contracts/stake-cw20-reward-distributor/schema/instantiate_msg.json
new file mode 100644
index 00000000..6d8a1ae1
--- /dev/null
+++ b/contracts/stake-cw20-reward-distributor/schema/instantiate_msg.json
@@ -0,0 +1,31 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "title": "InstantiateMsg",
+  "type": "object",
+  "required": [
+    "owner",
+    "reward_rate",
+    "reward_token",
+    "staking_addr"
+  ],
+  "properties": {
+    "owner": {
+      "type": "string"
+    },
+    "reward_rate": {
+      "$ref": "#/definitions/Uint128"
+    },
+    "reward_token": {
+      "type": "string"
+    },
+    "staking_addr": {
+      "type": "string"
+    }
+  },
+  "definitions": {
+    "Uint128": {
+      "description": "A thin wrapper around u128 that is using strings for JSON encoding/decoding, such that the full u128 range can be used for clients that convert JSON numbers to floats, like JavaScript and jq.\n\n# Examples\n\nUse `from` to create instances of this and `u128` to get the value out:\n\n``` # use cosmwasm_std::Uint128; let a = Uint128::from(123u128); assert_eq!(a.u128(), 123);\n\nlet b = Uint128::from(42u64); assert_eq!(b.u128(), 42);\n\nlet c = Uint128::from(70u32); assert_eq!(c.u128(), 70); ```",
+      "type": "string"
+    }
+  }
+}
diff --git a/contracts/stake-cw20-reward-distributor/schema/query_msg.json b/contracts/stake-cw20-reward-distributor/schema/query_msg.json
new file mode 100644
index 00000000..586c2ba9
--- /dev/null
+++ b/contracts/stake-cw20-reward-distributor/schema/query_msg.json
@@ -0,0 +1,18 @@
+{
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "title": "QueryMsg",
+  "oneOf": [
+    {
+      "type": "object",
+      "required": [
+        "info"
+      ],
+      "properties": {
+        "info": {
+          "type": "object"
+        }
+      },
+      "additionalProperties": false
+    }
+  ]
+}
diff --git a/contracts/stake-cw20-reward-distributor/src/contract.rs b/contracts/stake-cw20-reward-distributor/src/contract.rs
new file mode 100644
index 00000000..5f0b73db
--- /dev/null
+++ b/contracts/stake-cw20-reward-distributor/src/contract.rs
@@ -0,0 +1,178 @@
+use std::cmp::min;
+
+#[cfg(not(feature = "library"))]
+use cosmwasm_std::entry_point;
+use cosmwasm_std::{to_binary, Addr, CosmosMsg, StdError, Uint128, WasmMsg};
+
+use crate::error::ContractError;
+use crate::msg::{ExecuteMsg, InfoResponse, InstantiateMsg, QueryMsg};
+use crate::state::{Config, CONFIG, LAST_PAYMENT_BLOCK};
+use cosmwasm_std::{Binary, Deps, DepsMut, Env, MessageInfo, Response, StdResult};
+use cw2::set_contract_version;
+
+const CONTRACT_NAME: &str = "crates.io:stake-cw20-reward-distributor";
+const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");
+
+#[cfg_attr(not(feature = "library"), entry_point)]
+pub fn instantiate(
+    deps: DepsMut,
+    env: Env,
+    _info: MessageInfo,
+    msg: InstantiateMsg,
+) -> Result<Response, ContractError> {
+    set_contract_version(deps.storage, CONTRACT_NAME, CONTRACT_VERSION)?;
+
+    let owner = deps.api.addr_validate(&msg.owner)?;
+    let staking_addr = deps.api.addr_validate(&msg.staking_addr)?;
+    if !validate_staking(deps.as_ref(), staking_addr.clone()) {
+        return Err(ContractError::InvalidStakingContract {});
+    }
+
+    let reward_token = deps.api.addr_validate(&msg.reward_token)?;
+    if !validate_cw20(deps.as_ref(), reward_token.clone()) {
+        return Err(ContractError::InvalidCw20 {});
+    }
+
+    let config = Config {
+        owner,
+        staking_addr,
+        reward_token,
+        reward_rate: msg.reward_rate,
+    };
+    CONFIG.save(deps.storage, &config)?;
+
+    // Initialize last payment block
+    LAST_PAYMENT_BLOCK.save(deps.storage, &env.block.height)?;
+
+    Ok(Response::new().add_attribute("action", "instantiate"))
+}
+
+#[cfg_attr(not(feature = "library"), entry_point)]
+pub fn execute(
+    deps: DepsMut,
+    env: Env,
+    info: MessageInfo,
+    msg: ExecuteMsg,
+) -> Result<Response, ContractError> {
+    match msg {
+        ExecuteMsg::UpdateConfig {
+            owner,
+            staking_addr,
+            reward_rate,
+            reward_token,
+        } => execute_update_config(deps, info, owner, staking_addr, reward_rate, reward_token),
+        ExecuteMsg::Distribute {} => execute_distribute(deps, env),
+    }
+}
+
+pub fn execute_update_config(
+    deps: DepsMut,
+    info: MessageInfo,
+    owner: String,
+    staking_addr: String,
+    reward_rate: Uint128,
+    reward_token: String,
+) -> Result<Response, ContractError> {
+    let config = CONFIG.load(deps.storage)?;
+    if config.owner != info.sender {
+        return Err(ContractError::Unauthorized {});
+    }
+
+    let owner = deps.api.addr_validate(&owner)?;
+    let staking_addr = deps.api.addr_validate(&staking_addr)?;
+    if !validate_staking(deps.as_ref(), staking_addr.clone()) {
+        return Err(ContractError::InvalidStakingContract {});
+    }
+
+    let reward_token = deps.api.addr_validate(&reward_token)?;
+    if !validate_cw20(deps.as_ref(), reward_token.clone()) {
+        return Err(ContractError::InvalidCw20 {});
+    }
+
+    let config = Config {
+        owner,
+        staking_addr,
+        reward_token,
+        reward_rate,
+    };
+    CONFIG.save(deps.storage, &config)?;
+    Ok(Response::default())
+}
+
+pub fn validate_cw20(deps: Deps, cw20_addr: Addr) -> bool {
+    let response: Result<cw20::TokenInfoResponse, StdError> = deps
+        .querier
+        .query_wasm_smart(cw20_addr, &cw20::Cw20QueryMsg::TokenInfo {});
+    response.is_ok()
+}
+
+pub fn validate_staking(deps: Deps, staking_addr: Addr) -> bool {
+    let response: Result<stake_cw20::msg::TotalStakedAtHeightResponse, StdError> =
+        deps.querier.query_wasm_smart(
+            staking_addr,
+            &stake_cw20::msg::QueryMsg::TotalStakedAtHeight { height: None },
+        );
+    response.is_ok()
+}
+
+pub fn execute_distribute(deps: DepsMut, env: Env) -> Result<Response, ContractError> {
+    let config = CONFIG.load(deps.storage)?;
+    let last_payment_block = LAST_PAYMENT_BLOCK.load(deps.storage)?;
+    let block_diff = env.block.height - last_payment_block;
+    let pending_rewards: Uint128 = config.reward_rate * Uint128::new(block_diff.into());
+
+    if pending_rewards == Uint128::zero() {
+        return Err(ContractError::NoPendingPayments {});
+    }
+
+    let balance_info: cw20::BalanceResponse = deps.querier.query_wasm_smart(
+        config.reward_token.clone(),
+        &cw20::Cw20QueryMsg::Balance {
+            address: env.contract.address.to_string(),
+        },
+    )?;
+
+    if balance_info.balance == Uint128::zero() {
+        return Err(ContractError::OutOfFunds {});
+    }
+
+    let amount = min(balance_info.balance, pending_rewards);
+    LAST_PAYMENT_BLOCK.save(deps.storage, &env.block.height)?;
+
+    let msg = to_binary(&cw20::Cw20ExecuteMsg::Send {
+        contract: config.staking_addr.clone().into_string(),
+        amount,
+        msg: to_binary(&stake_cw20::msg::ReceiveMsg::Fund {}).unwrap(),
+    })?;
+    let send_msg: CosmosMsg = WasmMsg::Execute {
+        contract_addr: config.reward_token.into(),
+        msg,
+        funds: vec![],
+    }
+    .into();
+    Ok(Response::default().add_message(send_msg))
+}
+
+#[cfg_attr(not(feature = "library"), entry_point)]
+pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
+    match msg {
+        QueryMsg::Info {} => to_binary(&query_info(deps, env)?),
+    }
+}
+
+fn query_info(deps: Deps, env: Env) -> StdResult<InfoResponse> {
+    let config = CONFIG.load(deps.storage)?;
+    let last_payment_block = LAST_PAYMENT_BLOCK.load(deps.storage)?;
+    let balance_info: cw20::BalanceResponse = deps.querier.query_wasm_smart(
+        config.reward_token.clone(),
+        &cw20::Cw20QueryMsg::Balance {
+            address: env.contract.address.to_string(),
+        },
+    )?;
+
+    Ok(InfoResponse {
+        config,
+        last_payment_block,
+        balance: balance_info.balance,
+    })
+}
diff --git a/contracts/stake-cw20-reward-distributor/src/error.rs b/contracts/stake-cw20-reward-distributor/src/error.rs
new file mode 100644
index 00000000..198dc1de
--- /dev/null
+++ b/contracts/stake-cw20-reward-distributor/src/error.rs
@@ -0,0 +1,23 @@
+use cosmwasm_std::StdError;
+use thiserror::Error;
+
+#[derive(Error, Debug, PartialEq)]
+pub enum ContractError {
+    #[error("{0}")]
+    Std(#[from] StdError),
+
+    #[error("Unauthorized")]
+    Unauthorized {},
+
+    #[error("No Pending Payments")]
+    NoPendingPayments {},
+
+    #[error("Out of Funds")]
+    OutOfFunds {},
+
+    #[error("Invalid Cw20")]
+    InvalidCw20 {},
+
+    #[error("Invalid Staking Contract")]
+    InvalidStakingContract {},
+}
diff --git a/contracts/stake-cw20-reward-distributor/src/lib.rs b/contracts/stake-cw20-reward-distributor/src/lib.rs
new file mode 100644
index 00000000..8e58a2f1
--- /dev/null
+++ b/contracts/stake-cw20-reward-distributor/src/lib.rs
@@ -0,0 +1,9 @@
+pub mod contract;
+mod error;
+pub mod msg;
+pub mod state;
+
+#[cfg(test)]
+mod tests;
+
+pub use crate::error::ContractError;
diff --git a/contracts/stake-cw20-reward-distributor/src/msg.rs b/contracts/stake-cw20-reward-distributor/src/msg.rs
new file mode 100644
index 00000000..4e8e7e88
--- /dev/null
+++ b/contracts/stake-cw20-reward-distributor/src/msg.rs
@@ -0,0 +1,37 @@
+use crate::state::Config;
+use cosmwasm_std::Uint128;
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
+pub struct InstantiateMsg {
+    pub owner: String,
+    pub staking_addr: String,
+    pub reward_rate: Uint128,
+    pub reward_token: String,
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum ExecuteMsg {
+    UpdateConfig {
+        owner: String,
+        staking_addr: String,
+        reward_rate: Uint128,
+        reward_token: String,
+    },
+    Distribute {},
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum QueryMsg {
+    Info {},
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
+pub struct InfoResponse {
+    pub config: Config,
+    pub last_payment_block: u64,
+    pub balance: Uint128,
+}
diff --git a/contracts/stake-cw20-reward-distributor/src/state.rs b/contracts/stake-cw20-reward-distributor/src/state.rs
new file mode 100644
index 00000000..b1ad41b4
--- /dev/null
+++ b/contracts/stake-cw20-reward-distributor/src/state.rs
@@ -0,0 +1,16 @@
+use cosmwasm_std::{Addr, Uint128};
+use cw_storage_plus::Item;
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
+pub struct Config {
+    pub owner: Addr,
+    pub staking_addr: Addr,
+    pub reward_rate: Uint128,
+    pub reward_token: Addr,
+}
+
+pub const CONFIG: Item<Config> = Item::new("config");
+
+pub const LAST_PAYMENT_BLOCK: Item<u64> = Item::new("last_payment_block");


Test patch:
diff --git a/contracts/stake-cw20-reward-distributor/src/tests.rs b/contracts/stake-cw20-reward-distributor/src/tests.rs
new file mode 100644
index 00000000..738ff9cd
--- /dev/null
+++ b/contracts/stake-cw20-reward-distributor/src/tests.rs
@@ -0,0 +1,371 @@
+use crate::{
+    msg::{ExecuteMsg, InfoResponse, InstantiateMsg, QueryMsg},
+    state::Config,
+    ContractError,
+};
+use cosmwasm_std::{Addr, Empty, Uint128};
+use cw20::Cw20Coin;
+use cw_multi_test::{App, Contract, ContractWrapper, Executor};
+
+const OWNER: &str = "owner";
+const OWNER2: &str = "owner2";
+const MANAGER: &str = "manager";
+
+pub fn cw20_contract() -> Box<dyn Contract<Empty>> {
+    let contract = ContractWrapper::new(
+        cw20_base::contract::execute,
+        cw20_base::contract::instantiate,
+        cw20_base::contract::query,
+    );
+    Box::new(contract)
+}
+
+fn staking_contract() -> Box<dyn Contract<Empty>> {
+    let contract = ContractWrapper::new(
+        stake_cw20::contract::execute,
+        stake_cw20::contract::instantiate,
+        stake_cw20::contract::query,
+    );
+    Box::new(contract)
+}
+
+fn distributor_contract() -> Box<dyn Contract<Empty>> {
+    let contract = ContractWrapper::new(
+        crate::contract::execute,
+        crate::contract::instantiate,
+        crate::contract::query,
+    );
+    Box::new(contract)
+}
+
+fn instantiate_cw20(app: &mut App, initial_balances: Vec<Cw20Coin>) -> Addr {
+    let cw20_id = app.store_code(cw20_contract());
+    let msg = cw20_base::msg::InstantiateMsg {
+        name: String::from("Test"),
+        symbol: String::from("TEST"),
+        decimals: 6,
+        initial_balances,
+        mint: None,
+        marketing: None,
+    };
+
+    app.instantiate_contract(cw20_id, Addr::unchecked(OWNER), &msg, &[], "cw20", None)
+        .unwrap()
+}
+
+fn instantiate_staking(app: &mut App, cw20_addr: Addr) -> Addr {
+    let staking_id = app.store_code(staking_contract());
+    let msg = stake_cw20::msg::InstantiateMsg {
+        owner: Some(OWNER.to_string()),
+        manager: Some(MANAGER.to_string()),
+        token_address: cw20_addr.to_string(),
+        unstaking_duration: None,
+    };
+    app.instantiate_contract(
+        staking_id,
+        Addr::unchecked(OWNER),
+        &msg,
+        &[],
+        "staking",
+        None,
+    )
+    .unwrap()
+}
+
+fn instantiate_distributor(app: &mut App, msg: InstantiateMsg) -> Addr {
+    let code_id = app.store_code(distributor_contract());
+    app.instantiate_contract(
+        code_id,
+        Addr::unchecked(OWNER), // TODO: Remove this?
+        &msg,
+        &[],
+        "distributor",
+        None,
+    )
+    .unwrap()
+}
+
+fn get_balance_cw20<T: Into<String>, U: Into<String>>(
+    app: &App,
+    contract_addr: T,
+    address: U,
+) -> Uint128 {
+    let msg = cw20::Cw20QueryMsg::Balance {
+        address: address.into(),
+    };
+    let result: cw20::BalanceResponse = app.wrap().query_wasm_smart(contract_addr, &msg).unwrap();
+    result.balance
+}
+
+fn get_info<T: Into<String>>(app: &App, distributor_addr: T) -> InfoResponse {
+    let result: InfoResponse = app
+        .wrap()
+        .query_wasm_smart(distributor_addr, &QueryMsg::Info {})
+        .unwrap();
+    result
+}
+
+#[test]
+fn test_instantiate() {
+    let mut app = App::default();
+
+    let cw20_addr = instantiate_cw20(&mut app, vec![]);
+    let staking_addr = instantiate_staking(&mut app, cw20_addr.clone());
+
+    let msg = InstantiateMsg {
+        owner: OWNER.to_string(),
+        staking_addr: staking_addr.to_string(),
+        reward_rate: Uint128::new(1),
+        reward_token: cw20_addr.to_string(),
+    };
+
+    let distributor_addr = instantiate_distributor(&mut app, msg);
+    let response: InfoResponse = app
+        .wrap()
+        .query_wasm_smart(&distributor_addr, &QueryMsg::Info {})
+        .unwrap();
+
+    assert_eq!(
+        response.config,
+        Config {
+            owner: Addr::unchecked(OWNER),
+            staking_addr,
+            reward_rate: Uint128::new(1),
+            reward_token: cw20_addr,
+        }
+    );
+    assert_eq!(response.last_payment_block, app.block_info().height);
+}
+
+#[test]
+fn test_update_config() {
+    let mut app = App::default();
+
+    let cw20_addr = instantiate_cw20(&mut app, vec![]);
+    let staking_addr = instantiate_staking(&mut app, cw20_addr.clone());
+
+    let msg = InstantiateMsg {
+        owner: OWNER.to_string(),
+        staking_addr: staking_addr.to_string(),
+        reward_rate: Uint128::new(1),
+        reward_token: cw20_addr.to_string(),
+    };
+    let distributor_addr = instantiate_distributor(&mut app, msg);
+
+    let msg = ExecuteMsg::UpdateConfig {
+        owner: OWNER2.to_string(),
+        staking_addr: staking_addr.to_string(),
+        reward_rate: Uint128::new(5),
+        reward_token: cw20_addr.to_string(),
+    };
+
+    app.execute_contract(Addr::unchecked(OWNER), distributor_addr.clone(), &msg, &[])
+        .unwrap();
+
+    let response: InfoResponse = app
+        .wrap()
+        .query_wasm_smart(&distributor_addr, &QueryMsg::Info {})
+        .unwrap();
+
+    assert_eq!(
+        response.config,
+        Config {
+            owner: Addr::unchecked(OWNER2),
+            staking_addr: staking_addr.clone(),
+            reward_rate: Uint128::new(5),
+            reward_token: cw20_addr.clone(),
+        }
+    );
+
+    let msg = ExecuteMsg::UpdateConfig {
+        owner: OWNER2.to_string(),
+        staking_addr: staking_addr.to_string(),
+        reward_rate: Uint128::new(7),
+        reward_token: cw20_addr.to_string(),
+    };
+
+    let err: ContractError = app
+        .execute_contract(Addr::unchecked(OWNER), distributor_addr, &msg, &[])
+        .unwrap_err()
+        .downcast()
+        .unwrap();
+
+    assert_eq!(err, ContractError::Unauthorized {});
+}
+
+#[test]
+fn test_distribute() {
+    let mut app = App::default();
+
+    let cw20_addr = instantiate_cw20(
+        &mut app,
+        vec![cw20::Cw20Coin {
+            address: OWNER.to_string(),
+            amount: Uint128::from(1000u64),
+        }],
+    );
+    let staking_addr = instantiate_staking(&mut app, cw20_addr.clone());
+
+    let msg = InstantiateMsg {
+        owner: OWNER.to_string(),
+        staking_addr: staking_addr.to_string(),
+        reward_rate: Uint128::new(1),
+        reward_token: cw20_addr.to_string(),
+    };
+    let distributor_addr = instantiate_distributor(&mut app, msg);
+
+    let msg = cw20::Cw20ExecuteMsg::Transfer {
+        recipient: distributor_addr.to_string(),
+        amount: Uint128::from(1000u128),
+    };
+    app.execute_contract(Addr::unchecked(OWNER), cw20_addr.clone(), &msg, &[])
+        .unwrap();
+
+    app.update_block(|mut block| block.height += 10);
+    app.execute_contract(
+        Addr::unchecked(OWNER),
+        distributor_addr.clone(),
+        &ExecuteMsg::Distribute {},
+        &[],
+    )
+    .unwrap();
+
+    let staking_balance = get_balance_cw20(&app, cw20_addr.clone(), staking_addr.clone());
+    assert_eq!(staking_balance, Uint128::new(10));
+
+    let distributor_info = get_info(&app, distributor_addr.clone());
+    assert_eq!(distributor_info.balance, Uint128::new(990));
+    assert_eq!(distributor_info.last_payment_block, app.block_info().height);
+
+    app.update_block(|mut block| block.height += 500);
+    app.execute_contract(
+        Addr::unchecked(OWNER),
+        distributor_addr.clone(),
+        &ExecuteMsg::Distribute {},
+        &[],
+    )
+    .unwrap();
+
+    let staking_balance = get_balance_cw20(&app, cw20_addr.clone(), staking_addr.clone());
+    assert_eq!(staking_balance, Uint128::new(510));
+
+    let distributor_info = get_info(&app, distributor_addr.clone());
+    assert_eq!(distributor_info.balance, Uint128::new(490));
+    assert_eq!(distributor_info.last_payment_block, app.block_info().height);
+
+    app.update_block(|mut block| block.height += 1000);
+    app.execute_contract(
+        Addr::unchecked(OWNER),
+        distributor_addr.clone(),
+        &ExecuteMsg::Distribute {},
+        &[],
+    )
+    .unwrap();
+
+    let staking_balance = get_balance_cw20(&app, cw20_addr, staking_addr);
+    assert_eq!(staking_balance, Uint128::new(1000));
+
+    let distributor_info = get_info(&app, distributor_addr);
+    assert_eq!(distributor_info.balance, Uint128::new(0));
+    assert_eq!(distributor_info.last_payment_block, app.block_info().height);
+}
+
+#[test]
+fn test_instantiate_invalid_addrs() {
+    let mut app = App::default();
+    let cw20_addr = instantiate_cw20(
+        &mut app,
+        vec![cw20::Cw20Coin {
+            address: OWNER.to_string(),
+            amount: Uint128::from(1000u64),
+        }],
+    );
+    let staking_addr = instantiate_staking(&mut app, cw20_addr.clone());
+
+    let msg = InstantiateMsg {
+        owner: OWNER.to_string(),
+        staking_addr: staking_addr.to_string(),
+        reward_rate: Uint128::new(1),
+        reward_token: "invalid_cw20".to_string(),
+    };
+
+    let code_id = app.store_code(distributor_contract());
+    let err: ContractError = app
+        .instantiate_contract(
+            code_id,
+            Addr::unchecked(OWNER),
+            &msg,
+            &[],
+            "distributor",
+            None,
+        )
+        .unwrap_err()
+        .downcast()
+        .unwrap();
+
+    assert_eq!(err, ContractError::InvalidCw20 {});
+
+    let msg = InstantiateMsg {
+        owner: OWNER.to_string(),
+        staking_addr: "invalid_staking".to_string(),
+        reward_rate: Uint128::new(1),
+        reward_token: cw20_addr.to_string(),
+    };
+    let err: ContractError = app
+        .instantiate_contract(
+            code_id,
+            Addr::unchecked(OWNER),
+            &msg,
+            &[],
+            "distributor",
+            None,
+        )
+        .unwrap_err()
+        .downcast()
+        .unwrap();
+    assert_eq!(err, ContractError::InvalidStakingContract {});
+}
+
+#[test]
+fn test_update_config_invalid_addrs() {
+    let mut app = App::default();
+
+    let cw20_addr = instantiate_cw20(&mut app, vec![]);
+    let staking_addr = instantiate_staking(&mut app, cw20_addr.clone());
+
+    let msg = InstantiateMsg {
+        owner: OWNER.to_string(),
+        staking_addr: staking_addr.to_string(),
+        reward_rate: Uint128::new(1),
+        reward_token: cw20_addr.to_string(),
+    };
+    let distributor_addr = instantiate_distributor(&mut app, msg);
+
+    let msg = ExecuteMsg::UpdateConfig {
+        owner: OWNER.to_string(),
+        staking_addr: staking_addr.to_string(),
+        reward_rate: Uint128::new(5),
+        reward_token: "invalid_cw20".to_string(),
+    };
+
+    let err: ContractError = app
+        .execute_contract(Addr::unchecked(OWNER), distributor_addr.clone(), &msg, &[])
+        .unwrap_err()
+        .downcast()
+        .unwrap();
+    assert_eq!(err, ContractError::InvalidCw20 {});
+
+    let msg = ExecuteMsg::UpdateConfig {
+        owner: OWNER.to_string(),
+        staking_addr: "invalid_staking".to_string(),
+        reward_rate: Uint128::new(5),
+        reward_token: staking_addr.to_string(),
+    };
+
+    let err: ContractError = app
+        .execute_contract(Addr::unchecked(OWNER), distributor_addr, &msg, &[])
+        .unwrap_err()
+        .downcast()
+        .unwrap();
+    assert_eq!(err, ContractError::InvalidStakingContract {});
+}