Repository: CosmWasm/sylvia (rust)

Issue:
Support defining custom entry points
# Motivation

As described in #139 although user can define custom `entry_points` and use them in the generated contract it is not possible to use it in generated `multitests`.

# Solution

We could either use macro attribute pointing to the custom entry_point function or add method attribute inside of the macro.
Macro attribute approach will be more flexible for the user.

## Add missing Context type

```rust
pub struct SudoCtx<'a, C: CustomQuery = Empty> {
    pub deps: DepsMut<'a, C>,
    pub env: Env,
}
```

## Detect attribute `sv::custom_entry_point` on contract

This attribute should be used both by `contract` macro to use the custom entry point in `multitest helpers` and by `entry_points` macro so that there will be no duplication after generation.

```rust
pub mod entry_points {
    #[entry_point]
    pub fn sudo(_deps: DepsMut, _env: Env) -> ... {
    }
}

#[cfg_attr(not(feature = "library"), entry_points)]
#[contract]
#[sv::custom_entry_point(sudo=entry_points::sudo)]
impl ReceiverContract {
    pub const fn new() -> Self {
        Self {}
    }

    #[msg(instantiate)]
    pub fn instantiate(&self, _ctx: InstantiateCtx) -> StdResult<Response> {
        todo!()    
    }
}
``` 

## Allow overriding all entry points

This is the last step which aim is to make `sylvia` more flexible for the user. In case of some missing feature or individual need user should be able to override all of the entry points.

# Alternative

Although support for `sudo` is not yet implemented we could allow defining `sudo` msg as it is done in case of `reply`. If detected `sylvia` will generate appropriate `entry_point` and implement `cw_multi_test::Contract` with this `entry_point`.

Pull request:
feat: Entry points overriding
close #189 

Patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index fe3aaa5..5040791 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -84,6 +84,9 @@ jobs:
       - name: Build cw1-whitelist
         working-directory: examples/contracts/cw1-whitelist
         run: cargo build --release --target wasm32-unknown-unknown --locked --lib
+      - name: Build entry-points-overriding
+        working-directory: examples/contracts/entry-points-overriding
+        run: cargo build --release --target wasm32-unknown-unknown --locked --lib
       - name: Install cosmwasm-check
         run: cargo install cosmwasm-check --force
       - name: Check contracts
@@ -97,6 +100,9 @@ jobs:
       - name: Cw20-base schema
         working-directory: examples/contracts/cw20-base
         run: cargo schema
+      - name: Entry-points-overriding schema
+        working-directory: examples/contracts/entry-points-overriding
+        run: cargo schema
       - name: Cw1-whitelist ts-codegen 
         working-directory: examples/contracts/cw1-whitelist/
         run: cosmwasm-ts-codegen generate --plugin client --schema ./schema --out ./ts --name cw1-whitelist --no-bundle
@@ -106,6 +112,9 @@ jobs:
       - name: Cw20-base ts-codegen 
         working-directory: examples/contracts/cw20-base/
         run: cosmwasm-ts-codegen generate --plugin client --schema ./schema --out ./ts --name cw20-base --no-bundle
+      - name: Entry-points-overriding ts-codegen
+        working-directory: examples/contracts/entry-points-overriding
+        run: cosmwasm-ts-codegen generate --plugin client --schema ./schema --out ./ts --name entry-points-overriding --no-bundle
       - name: Archive schema artifats
         uses: actions/upload-artifact@v3
         with:
@@ -114,6 +123,7 @@ jobs:
             examples/contracts/cw1-subkeys/schema/cw1-subkeys.json
             examples/contracts/cw1-whitelist/schema/cw1-whitelist.json
             examples/contracts/cw20-base/schema/cw20-base.json
+            examples/contracts/entry-points-overriding/schema/entry-points-overriding.json
 
   coverage:
     name: Code coverage
diff --git a/Cargo.toml b/Cargo.toml
index 95cb17f..c7ee65d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,9 +1,5 @@
 [workspace]
-members = [
-  "sylvia",
-  "sylvia-derive",
-
-]
+members = ["sylvia", "sylvia-derive"]
 exclude = ["examples/*"]
 resolver = "2"
 
diff --git a/examples/Cargo.lock b/examples/Cargo.lock
index abe1d92..7d33f84 100644
--- a/examples/Cargo.lock
+++ b/examples/Cargo.lock
@@ -506,6 +506,20 @@ dependencies = [
  "zeroize",
 ]
 
+[[package]]
+name = "entry-points-overriding"
+version = "0.5.0"
+dependencies = [
+ "anyhow",
+ "cosmwasm-schema",
+ "cosmwasm-std",
+ "cw-multi-test",
+ "cw-storage-plus",
+ "cw-utils",
+ "serde",
+ "sylvia",
+]
+
 [[package]]
 name = "ff"
 version = "0.12.1"
diff --git a/examples/Cargo.toml b/examples/Cargo.toml
index 0fa44cc..7314a3b 100644
--- a/examples/Cargo.toml
+++ b/examples/Cargo.toml
@@ -11,10 +11,10 @@ members = [
   "contracts/cw1-whitelist",
   "contracts/cw1-subkeys",
   "contracts/cw20-base",
+  "contracts/entry-points-overriding",
 ]
 resolver = "2"
 
 [workspace.package]
 version = "0.5.0"
 edition = "2021"
-
diff --git a/examples/contracts/entry-points-overriding/.cargo/config b/examples/contracts/entry-points-overriding/.cargo/config
new file mode 100644
index 0000000..d8ab80f
--- /dev/null
+++ b/examples/contracts/entry-points-overriding/.cargo/config
@@ -0,0 +1,6 @@
+[alias]
+wasm = "build --release --target wasm32-unknown-unknown --lib"
+wasm-debug = "build --target wasm32-unknown-unknown --lib"
+unit-test = "test --lib"
+integration-test = "test --test integration"
+schema = "run --bin schema"
diff --git a/examples/contracts/entry-points-overriding/Cargo.toml b/examples/contracts/entry-points-overriding/Cargo.toml
new file mode 100644
index 0000000..91dbdf8
--- /dev/null
+++ b/examples/contracts/entry-points-overriding/Cargo.toml
@@ -0,0 +1,31 @@
+[package]
+name = "entry-points-overriding"
+version = { workspace = true }
+authors = ["Jan Woźniak <jan@confio.gmbh>"]
+edition = { workspace = true }
+description = "Example usage of sudo and entry point overriding."
+license = "Apache-2.0"
+repository = "https://github.com/CosmWasm/sylvia"
+homepage = "https://cosmwasm.com"
+
+[lib]
+crate-type = ["cdylib", "rlib"]
+
+[features]
+library = []
+tests = ["library", "cw-multi-test", "anyhow"]
+
+[dependencies]
+anyhow = { version = "1.0", optional = true }
+cosmwasm-schema = "1.2"
+cosmwasm-std = { version = "1.2", features = ["staking"] }
+cw-multi-test = { version = "0.16", optional = true }
+cw-storage-plus = "1.0"
+cw-utils = "1.0"
+serde = { version = "1.0", default-features = false, features = ["derive"] }
+sylvia = { path = "../../../sylvia" }
+
+[dev-dependencies]
+anyhow = "1.0"
+cw-multi-test = "0.16"
+sylvia = { path = "../../../sylvia", features = ["mt"] }
diff --git a/examples/contracts/entry-points-overriding/src/bin/schema.rs b/examples/contracts/entry-points-overriding/src/bin/schema.rs
new file mode 100644
index 0000000..0c9fdd4
--- /dev/null
+++ b/examples/contracts/entry-points-overriding/src/bin/schema.rs
@@ -0,0 +1,15 @@
+use cosmwasm_schema::write_api;
+
+use entry_points_overriding::contract::{ContractQueryMsg, InstantiateMsg};
+use entry_points_overriding::messages::CustomExecMsg;
+use entry_points_overriding::messages::SudoMsg;
+
+#[cfg(not(tarpaulin_include))]
+fn main() {
+    write_api! {
+        instantiate: InstantiateMsg,
+        execute: CustomExecMsg,
+        query: ContractQueryMsg,
+        sudo: SudoMsg
+    }
+}
diff --git a/examples/contracts/entry-points-overriding/src/contract.rs b/examples/contracts/entry-points-overriding/src/contract.rs
new file mode 100644
index 0000000..962891d
--- /dev/null
+++ b/examples/contracts/entry-points-overriding/src/contract.rs
@@ -0,0 +1,46 @@
+use cosmwasm_std::{Response, StdError, StdResult};
+use cw_storage_plus::Item;
+use sylvia::types::{ExecCtx, InstantiateCtx, QueryCtx};
+use sylvia::{contract, schemars};
+
+#[cfg(not(feature = "library"))]
+use sylvia::entry_points;
+
+use crate::messages::CountResponse;
+
+pub struct CounterContract {
+    pub(crate) counter: Item<'static, u32>,
+}
+
+#[cfg_attr(not(feature = "library"), entry_points)]
+#[contract]
+#[sv::override_entry_point(sudo=crate::entry_points::sudo(crate::messages::SudoMsg))]
+#[sv::override_entry_point(exec=crate::entry_points::execute(crate::messages::CustomExecMsg))]
+impl CounterContract {
+    pub const fn new() -> Self {
+        Self {
+            counter: Item::new("counter"),
+        }
+    }
+
+    #[msg(instantiate)]
+    pub fn instantiate(&self, ctx: InstantiateCtx) -> StdResult<Response> {
+        self.counter.save(ctx.deps.storage, &0)?;
+        Ok(Response::new())
+    }
+
+    #[msg(query)]
+    pub fn count(&self, ctx: QueryCtx) -> StdResult<CountResponse> {
+        let count = self.counter.load(ctx.deps.storage)?;
+        Ok(CountResponse { count })
+    }
+
+    #[msg(exec)]
+    pub fn increase_by_two(&self, ctx: ExecCtx) -> StdResult<Response> {
+        self.counter
+            .update(ctx.deps.storage, |count| -> Result<u32, StdError> {
+                Ok(count + 2)
+            })?;
+        Ok(Response::new())
+    }
+}
diff --git a/examples/contracts/entry-points-overriding/src/entry_points.rs b/examples/contracts/entry-points-overriding/src/entry_points.rs
new file mode 100644
index 0000000..6fc1746
--- /dev/null
+++ b/examples/contracts/entry-points-overriding/src/entry_points.rs
@@ -0,0 +1,20 @@
+use cosmwasm_std::{entry_point, DepsMut, Env, MessageInfo, Response, StdResult};
+
+use crate::contract::CounterContract;
+use crate::messages::{CustomExecMsg, SudoMsg};
+
+#[entry_point]
+pub fn sudo(deps: DepsMut, _env: Env, _msg: SudoMsg) -> StdResult<Response> {
+    CounterContract::new().counter.save(deps.storage, &3)?;
+    Ok(Response::new())
+}
+
+#[entry_point]
+pub fn execute(
+    deps: DepsMut,
+    env: Env,
+    info: MessageInfo,
+    msg: CustomExecMsg,
+) -> StdResult<Response> {
+    msg.dispatch((deps, env, info))
+}
diff --git a/examples/contracts/entry-points-overriding/src/lib.rs b/examples/contracts/entry-points-overriding/src/lib.rs
new file mode 100644
index 0000000..f8b45b5
--- /dev/null
+++ b/examples/contracts/entry-points-overriding/src/lib.rs
@@ -0,0 +1,5 @@
+pub mod contract;
+pub mod entry_points;
+pub mod messages;
+#[cfg(any(test, feature = "tests"))]
+pub mod multitest;
diff --git a/examples/contracts/entry-points-overriding/src/messages.rs b/examples/contracts/entry-points-overriding/src/messages.rs
new file mode 100644
index 0000000..c032e2b
--- /dev/null
+++ b/examples/contracts/entry-points-overriding/src/messages.rs
@@ -0,0 +1,44 @@
+use cosmwasm_schema::cw_serde;
+use cosmwasm_std::{DepsMut, Env, MessageInfo, Response, StdError, StdResult};
+use sylvia::types::ExecCtx;
+
+use crate::contract::{ContractExecMsg, CounterContract};
+
+#[cw_serde]
+pub struct CountResponse {
+    pub count: u32,
+}
+
+#[cw_serde]
+pub enum SudoMsg {
+    SetCountToThree {},
+}
+
+#[cw_serde]
+pub enum UserExecMsg {
+    IncreaseByOne {},
+}
+
+pub fn increase_by_one(ctx: ExecCtx) -> StdResult<Response> {
+    CounterContract::new()
+        .counter
+        .update(ctx.deps.storage, |count| -> Result<u32, StdError> {
+            Ok(count + 1)
+        })?;
+    Ok(Response::new())
+}
+
+#[cw_serde]
+pub enum CustomExecMsg {
+    ContractExec(ContractExecMsg),
+    CustomExec(UserExecMsg),
+}
+
+impl CustomExecMsg {
+    pub fn dispatch(self, ctx: (DepsMut, Env, MessageInfo)) -> StdResult<Response> {
+        match self {
+            CustomExecMsg::ContractExec(msg) => msg.dispatch(&CounterContract::new(), ctx),
+            CustomExecMsg::CustomExec(_) => increase_by_one(ctx.into()),
+        }
+    }
+}
diff --git a/sylvia-derive/src/input.rs b/sylvia-derive/src/input.rs
index 34b8643..87f8deb 100644
--- a/sylvia-derive/src/input.rs
+++ b/sylvia-derive/src/input.rs
@@ -9,7 +9,7 @@ use crate::crate_module;
 use crate::interfaces::Interfaces;
 use crate::message::{ContractEnumMessage, EnumMessage, GlueMessage, MsgVariants, StructMessage};
 use crate::multitest::{MultitestHelpers, TraitMultitestHelpers};
-use crate::parser::{ContractArgs, ContractErrorAttr, Custom, MsgType};
+use crate::parser::{ContractArgs, ContractErrorAttr, Custom, MsgType, OverrideEntryPoints};
 use crate::remote::Remote;
 use crate::variant_descs::AsVariantDescs;
 
@@ -27,6 +27,7 @@ pub struct ImplInput<'a> {
     item: &'a ItemImpl,
     generics: Vec<&'a GenericParam>,
     custom: Custom<'a>,
+    override_entry_points: OverrideEntryPoints,
 }
 
 impl<'a> TraitInput<'a> {
@@ -126,6 +127,7 @@ impl<'a> ImplInput<'a> {
             .unwrap_or_else(|| parse_quote! { #sylvia ::cw_std::StdError });
 
         let custom = Custom::new(&item.attrs);
+        let override_entry_points = OverrideEntryPoints::new(&item.attrs);
 
         Self {
             attributes,
@@ -133,6 +135,7 @@ impl<'a> ImplInput<'a> {
             generics,
             error,
             custom,
+            override_entry_points,
         }
     }
 
@@ -145,6 +148,7 @@ impl<'a> ImplInput<'a> {
                 &self.error,
                 &self.generics,
                 &self.custom,
+                &self.override_entry_points,
             )
             .emit()
         } else {
diff --git a/sylvia-derive/src/message.rs b/sylvia-derive/src/message.rs
index a1dfef0..f57f592 100644
--- a/sylvia-derive/src/message.rs
+++ b/sylvia-derive/src/message.rs
@@ -2,6 +2,7 @@ use crate::check_generics::CheckGenerics;
 use crate::crate_module;
 use crate::parser::{
     parse_struct_message, ContractErrorAttr, ContractMessageAttr, Custom, MsgAttr, MsgType,
+    OverrideEntryPoint, OverrideEntryPoints,
 };
 use crate::strip_generics::StripGenerics;
 use crate::utils::{extract_return_type, filter_wheres, process_fields};
@@ -544,7 +545,7 @@ impl<'a> MsgVariant<'a> {
     /// Emits match leg dispatching against this variant. Assumes enum variants are imported into the
     /// scope. Dispatching is performed by calling the function this variant is build from on the
     /// `contract` variable, with `ctx` as its first argument - both of them should be in scope.
-    pub fn emit_dispatch_leg(&self, msg_attr: MsgType) -> TokenStream {
+    pub fn emit_dispatch_leg(&self, msg_type: MsgType) -> TokenStream {
         use MsgType::*;
 
         let Self {
@@ -568,7 +569,7 @@ impl<'a> MsgVariant<'a> {
             .map(|(field, num_field)| quote!(#field : #num_field));
 
         #[cfg(not(tarpaulin_include))]
-        match msg_attr {
+        match msg_type {
             Exec => quote! {
                 #name {
                     #(#fields,)*
@@ -579,8 +580,8 @@ impl<'a> MsgVariant<'a> {
                     #(#fields,)*
                 } => #sylvia ::cw_std::to_binary(&contract.#function_name(Into::into(ctx), #(#args),*)?).map_err(Into::into)
             },
-            Instantiate | Migrate | Reply => {
-                emit_error!(name.span(), "Instantiation, Reply and Migrate messages not supported on traits, they should be defined on contracts directly");
+            Instantiate | Migrate | Reply | Sudo => {
+                emit_error!(name.span(), "Instantiation, Reply, Migrate and Sudo messages not supported on traits, they should be defined on contracts directly");
                 quote! {}
             }
         }
@@ -1093,12 +1094,14 @@ pub struct EntryPoints<'a> {
     error: Type,
     reply: Option<Ident>,
     custom: Custom<'a>,
+    override_entry_points: OverrideEntryPoints,
 }
 
 impl<'a> EntryPoints<'a> {
     pub fn new(source: &'a ItemImpl) -> Self {
         let sylvia = crate_module();
         let name = StripGenerics.fold_type(*source.self_ty.clone());
+        let override_entry_points = OverrideEntryPoints::new(&source.attrs);
 
         let error = source
             .attrs
@@ -1128,6 +1131,7 @@ impl<'a> EntryPoints<'a> {
             error,
             reply,
             custom,
+            override_entry_points,
         }
     }
 
@@ -1137,59 +1141,52 @@ impl<'a> EntryPoints<'a> {
             error,
             reply,
             custom,
+            override_entry_points,
         } = self;
         let sylvia = crate_module();
 
         let custom_msg = custom.msg();
-        #[cfg(not(tarpaulin_include))]
-        let resp_type = quote! { #sylvia ::cw_std::Response < #custom_msg > };
 
         #[cfg(not(tarpaulin_include))]
         {
-            let reply = match reply {
-                Some(reply) => quote! {
-                    #[#sylvia ::cw_std::entry_point]
-                    pub fn reply(
-                        deps: #sylvia ::cw_std::DepsMut,
-                        env: #sylvia ::cw_std::Env,
-                        msg: #sylvia ::cw_std::Reply,
-                    ) -> Result<#resp_type, #error> {
-                        #name ::new(). #reply((deps, env).into(), msg).map_err(Into::into)
-                    }
-                },
-                None => quote! {},
-            };
+            let entry_points = [MsgType::Instantiate, MsgType::Exec, MsgType::Query]
+                .into_iter()
+                .map(
+                    |msg_type| match override_entry_points.get_entry_point(msg_type) {
+                        Some(_) => quote! {},
+                        None => OverrideEntryPoint::emit_default_entry_point(
+                            &custom_msg,
+                            name,
+                            error,
+                            msg_type,
+                        ),
+                    },
+                );
+
+            let reply_ep = override_entry_points
+                .get_entry_point(MsgType::Reply)
+                .map(|_| quote! {})
+                .unwrap_or_else(|| match reply {
+                    Some(reply) => quote! {
+                        #[#sylvia ::cw_std::entry_point]
+                        pub fn reply(
+                            deps: #sylvia ::cw_std::DepsMut,
+                            env: #sylvia ::cw_std::Env,
+                            msg: #sylvia ::cw_std::Reply,
+                        ) -> Result<#sylvia ::cw_std::Response < #custom_msg >, #error> {
+                            #name ::new(). #reply((deps, env).into(), msg).map_err(Into::into)
+                        }
+                    },
+                    _ => quote! {},
+                });
 
             quote! {
                 pub mod entry_points {
                     use super::*;
 
-                    #[#sylvia ::cw_std::entry_point]
-                    pub fn instantiate(
-                        deps: #sylvia ::cw_std::DepsMut,
-                        env: #sylvia ::cw_std::Env,
-                        info: #sylvia ::cw_std::MessageInfo,
-                        msg: InstantiateMsg,
-                    ) -> Result<#resp_type, #error> {
-                        msg.dispatch(&#name ::new() , (deps, env, info)).map_err(Into::into)
-                    }
-
-                    #[#sylvia ::cw_std::entry_point]
-                    pub fn execute(
-                        deps: #sylvia ::cw_std::DepsMut,
-                        env: #sylvia ::cw_std::Env,
-                        info: #sylvia ::cw_std::MessageInfo,
-                        msg: ContractExecMsg,
-                    ) -> Result<#resp_type, #error> {
-                        msg.dispatch(&#name ::new(), (deps, env, info)).map_err(Into::into)
-                    }
-
-                    #[#sylvia ::cw_std::entry_point]
-                    pub fn query(deps: #sylvia ::cw_std::Deps, env: #sylvia ::cw_std::Env, msg: ContractQueryMsg) -> Result<#sylvia ::cw_std::Binary, #error> {
-                        msg.dispatch(&#name ::new(), (deps, env)).map_err(Into::into)
-                    }
+                    #(#entry_points)*
 
-                    #reply
+                    #reply_ep
                 }
             }
         }
diff --git a/sylvia-derive/src/parser.rs b/sylvia-derive/src/parser.rs
index 0b145be..c64b96b 100644
--- a/sylvia-derive/src/parser.rs
+++ b/sylvia-derive/src/parser.rs
@@ -51,6 +51,7 @@ pub enum MsgType {
     Instantiate,
     Migrate,
     Reply,
+    Sudo,
 }
 
 /// `#[msg(...)]` attribute for `interface` macro
@@ -72,7 +73,7 @@ impl MsgType {
             Exec | Instantiate => quote! {
                 (#sylvia ::cw_std::DepsMut, #sylvia ::cw_std::Env, #sylvia ::cw_std::MessageInfo)
             },
-            Migrate | Reply => quote! {
+            Migrate | Reply | Sudo => quote! {
                 (#sylvia ::cw_std::DepsMut, #sylvia ::cw_std::Env)
             },
             Query => quote! {
@@ -81,6 +82,46 @@ impl MsgType {
         }
     }
 
+    pub fn emit_ctx_params(self) -> TokenStream {
+        use MsgType::*;
+
+        let sylvia = crate_module();
+
+        match self {
+            Exec | Instantiate => quote! {
+                deps: #sylvia ::cw_std::DepsMut, env: #sylvia ::cw_std::Env, info: #sylvia ::cw_std::MessageInfo
+            },
+            Migrate | Reply | Sudo => quote! {
+                deps: #sylvia ::cw_std::DepsMut, env: #sylvia ::cw_std::Env
+            },
+            Query => quote! {
+                deps: #sylvia ::cw_std::Deps, env: #sylvia ::cw_std::Env
+            },
+        }
+    }
+
+    pub fn emit_ep_name(self) -> TokenStream {
+        use MsgType::*;
+
+        match self {
+            Exec => quote! { execute },
+            Instantiate => quote! { instantiate },
+            Migrate => quote! { migrate },
+            Sudo => quote! { sudo },
+            Reply => quote! { reply },
+            Query => quote! { query },
+        }
+    }
+
+    pub fn emit_ctx_values(self) -> TokenStream {
+        use MsgType::*;
+
+        match self {
+            Exec | Instantiate => quote! { deps, env, info },
+            Migrate | Reply | Query | Sudo => quote! { deps, env },
+        }
+    }
+
     /// Emits type which should be returned by dispatch function for this kind of message
     pub fn emit_result_type(self, msg_type: &Type, err_type: &Type) -> TokenStream {
         use MsgType::*;
@@ -88,7 +129,7 @@ impl MsgType {
         let sylvia = crate_module();
 
         match self {
-            Exec | Instantiate | Migrate | Reply => {
+            Exec | Instantiate | Migrate | Reply | Sudo => {
                 quote! {
                     std::result::Result< #sylvia:: cw_std::Response <#msg_type>, #err_type>
                 }
@@ -98,6 +139,17 @@ impl MsgType {
             },
         }
     }
+
+    pub fn emit_msg_name(&self) -> Type {
+        match self {
+            MsgType::Exec => parse_quote! { ContractExecMsg },
+            MsgType::Query => parse_quote! { ContractQueryMsg },
+            MsgType::Instantiate => parse_quote! { InstantiateMsg },
+            MsgType::Migrate => parse_quote! { MigrateMsg },
+            MsgType::Reply => parse_quote! { ReplyMsg },
+            MsgType::Sudo => todo!(),
+        }
+    }
 }
 
 impl PartialEq<MsgType> for MsgAttr {
@@ -400,6 +452,142 @@ impl Parse for Custom<'_> {
     }
 }
 
+#[derive(Debug)]
+pub struct OverrideEntryPoint {
+    entry_point: Path,
+    msg_name: Type,
+    msg_type: MsgType,
+}
+
+impl OverrideEntryPoint {
+    pub fn emit_multitest_dispatch(&self) -> TokenStream {
+        let Self {
+            entry_point,
+            msg_name,
+            msg_type,
+        } = self;
+
+        let sylvia = crate_module();
+        let values = msg_type.emit_ctx_values();
+
+        quote! {
+            #entry_point ( #values .into(), #sylvia ::cw_std::from_slice::< #msg_name >(&msg)?)
+                .map_err(Into::into)
+        }
+    }
+
+    pub fn emit_multitest_default_dispatch(ty: MsgType) -> TokenStream {
+        let sylvia = crate_module();
+
+        let values = ty.emit_ctx_values();
+        let msg_name = ty.emit_msg_name();
+
+        quote! {
+            #sylvia ::cw_std::from_slice::< #msg_name >(&msg)?
+                .dispatch(self, ( #values ))
+                .map_err(Into::into)
+        }
+    }
+
+    pub fn emit_default_entry_point(
+        custom_msg: &Type,
+        name: &Type,
+        error: &Type,
+        msg_type: MsgType,
+    ) -> TokenStream {
+        let sylvia = crate_module();
+
+        #[cfg(not(tarpaulin_include))]
+        let resp_type = match msg_type {
+            MsgType::Query => quote! { #sylvia ::cw_std::Binary },
+            _ => quote! { #sylvia ::cw_std::Response < #custom_msg > },
+        };
+        let params = msg_type.emit_ctx_params();
+        let values = msg_type.emit_ctx_values();
+        let ep_name = msg_type.emit_ep_name();
+        let msg_name = msg_type.emit_msg_name();
+
+        quote! {
+            #[#sylvia ::cw_std::entry_point]
+            pub fn #ep_name (
+                #params ,
+                msg: #msg_name,
+            ) -> Result<#resp_type, #error> {
+                msg.dispatch(&#name ::new() , ( #values )).map_err(Into::into)
+            }
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct OverrideEntryPoints(Vec<OverrideEntryPoint>);
+
+impl OverrideEntryPoints {
+    pub fn new(attrs: &[Attribute]) -> Self {
+        let entry_points = attrs
+            .iter()
+            .filter(|attr| match sylvia_attribute(attr) {
+                Some(attr) => attr == "override_entry_point",
+                None => false,
+            })
+            .filter_map(
+                |attr| match OverrideEntryPoint::parse.parse2(attr.tokens.clone()) {
+                    Ok(entry_point) => Some(entry_point),
+                    Err(err) => {
+                        emit_error!(attr.span(), err);
+                        None
+                    }
+                },
+            )
+            .collect();
+
+        Self(entry_points)
+    }
+
+    pub fn get_entry_point(&self, ty: MsgType) -> Option<&OverrideEntryPoint> {
+        self.0.iter().find(|entry_point| entry_point.msg_type == ty)
+    }
+}
+
+#[cfg(not(tarpaulin_include))]
+// False negative. It is being called in closure
+impl Parse for OverrideEntryPoint {
+    fn parse(input: ParseStream) -> Result<Self> {
+        let content;
+        parenthesized!(content in input);
+
+        let ty: Ident = content.parse()?;
+        let _: Token![=] = content.parse()?;
+        let entry_point = content.parse()?;
+
+        let msg_content;
+        parenthesized!(msg_content in content);
+
+        let msg_name = msg_content.parse()?;
+
+        let msg_type = match ty.to_string().as_str() {
+            "exec" =>  MsgType::Exec,
+            "instantiate" =>  MsgType::Instantiate,
+            "query" =>  MsgType::Instantiate,
+            "migrate" => MsgType::Migrate,
+            "reply" => MsgType::Reply,
+            "sudo" =>  MsgType::Sudo,
+            &_ => {
+                return Err(Error::new(
+                    ty.span(),
+                    "Invalid entry point. Expected exec, instantiate, query, migrate, reply or sudo. Found {ty}",
+                ))
+            }
+        };
+
+        Ok(Self {
+            entry_point,
+            msg_name,
+            msg_type,
+        })
+    }
+}
+
 pub fn sylvia_attribute(attr: &Attribute) -> Option<&Ident> {
     if attr.path.segments.len() == 2 && attr.path.segments[0].ident == "sv" {
         Some(&attr.path.segments[1].ident)
diff --git a/sylvia/src/types.rs b/sylvia/src/types.rs
index 9b73180..133fd8c 100644
--- a/sylvia/src/types.rs
+++ b/sylvia/src/types.rs
@@ -27,6 +27,11 @@ pub struct QueryCtx<'a, C: CustomQuery = Empty> {
     pub env: Env,
 }
 
+pub struct SudoCtx<'a, C: CustomQuery = Empty> {
+    pub deps: DepsMut<'a, C>,
+    pub env: Env,
+}
+
 impl<C: CustomQuery> ExecCtx<'_, C> {
     pub fn branch(&'_ mut self) -> ExecCtx<'_, C> {
         ExecCtx {
@@ -36,6 +41,7 @@ impl<C: CustomQuery> ExecCtx<'_, C> {
         }
     }
 }
+
 impl<C: CustomQuery> InstantiateCtx<'_, C> {
     pub fn branch(&'_ mut self) -> InstantiateCtx<'_, C> {
         InstantiateCtx {
@@ -46,6 +52,15 @@ impl<C: CustomQuery> InstantiateCtx<'_, C> {
     }
 }
 
+impl<C: CustomQuery> SudoCtx<'_, C> {
+    pub fn branch(&'_ mut self) -> SudoCtx<'_, C> {
+        SudoCtx {
+            deps: self.deps.branch(),
+            env: self.env.clone(),
+        }
+    }
+}
+
 impl<'a, C: CustomQuery> From<(DepsMut<'a, C>, Env)> for MigrateCtx<'a, C> {
     fn from((deps, env): (DepsMut<'a, C>, Env)) -> Self {
         Self { deps, env }


Test patch:
diff --git a/examples/contracts/entry-points-overriding/src/multitest.rs b/examples/contracts/entry-points-overriding/src/multitest.rs
new file mode 100644
index 0000000..d822451
--- /dev/null
+++ b/examples/contracts/entry-points-overriding/src/multitest.rs
@@ -0,0 +1,71 @@
+#[cfg(test)]
+mod test {
+    use crate::{
+        contract::{multitest_utils::CodeId, ContractExecMsg, ExecMsg},
+        messages::{CustomExecMsg, SudoMsg, UserExecMsg},
+    };
+    use cosmwasm_std::Addr;
+    use cw_multi_test::Executor;
+    use sylvia::multitest::App;
+
+    #[test]
+    fn overriden_entry_points_in_mt() {
+        let app = App::default();
+        let code_id = CodeId::store_code(&app);
+
+        let owner = "owner";
+
+        let contract = code_id
+            .instantiate()
+            .with_label("Contract")
+            .with_admin(Some(owner))
+            .call(owner)
+            .unwrap();
+
+        let count = contract.count().unwrap().count;
+        assert_eq!(count, 0);
+
+        let msg = SudoMsg::SetCountToThree {};
+
+        contract
+            .app
+            .app_mut()
+            .wasm_sudo(contract.contract_addr.clone(), &msg)
+            .unwrap();
+
+        let count = contract.count().unwrap().count;
+        assert_eq!(count, 3);
+
+        // custom ExecMsg
+        let msg = CustomExecMsg::CustomExec(UserExecMsg::IncreaseByOne {});
+        (*contract.app)
+            .app_mut()
+            .execute_contract(
+                Addr::unchecked(owner),
+                contract.contract_addr.clone(),
+                &msg,
+                &[],
+            )
+            .unwrap();
+
+        let count = contract.count().unwrap().count;
+        assert_eq!(count, 4);
+
+        // custom ExecMsg
+        let msg = CustomExecMsg::ContractExec(ContractExecMsg::CounterContract(
+            ExecMsg::increase_by_two(),
+        ));
+        (*contract.app)
+            .app_mut()
+            .execute_contract(
+                Addr::unchecked(owner),
+                contract.contract_addr.clone(),
+                &msg,
+                &[],
+            )
+            .unwrap();
+
+        let count = contract.count().unwrap().count;
+        assert_eq!(count, 6);
+    }
+}
diff --git a/sylvia-derive/src/multitest.rs b/sylvia-derive/src/multitest.rs
index 3bee29d..ab077fd 100644
--- a/sylvia-derive/src/multitest.rs
+++ b/sylvia-derive/src/multitest.rs
@@ -9,7 +9,10 @@ use syn::{FnArg, GenericParam, ImplItem, ItemImpl, ItemTrait, Pat, PatType, Path
 use crate::check_generics::CheckGenerics;
 use crate::crate_module;
 use crate::message::MsgField;
-use crate::parser::{parse_struct_message, ContractMessageAttr, Custom, MsgAttr, MsgType};
+use crate::parser::{
+    parse_struct_message, ContractMessageAttr, Custom, MsgAttr, MsgType, OverrideEntryPoint,
+    OverrideEntryPoints,
+};
 use crate::utils::{extract_return_type, process_fields};
 
 struct MessageSignature<'a> {
@@ -32,6 +35,7 @@ pub struct MultitestHelpers<'a> {
     contract_name: &'a Ident,
     proxy_name: Ident,
     custom: &'a Custom<'a>,
+    override_entry_points: &'a OverrideEntryPoints,
 }
 
 fn interface_name(source: &ItemImpl) -> &Ident {
@@ -60,6 +64,7 @@ impl<'a> MultitestHelpers<'a> {
         contract_error: &'a Type,
         generics: &'a [&'a GenericParam],
         custom: &'a Custom,
+        override_entry_points: &'a OverrideEntryPoints,
     ) -> Self {
         let mut is_migrate = false;
         let mut reply = None;
@@ -197,6 +202,7 @@ impl<'a> MultitestHelpers<'a> {
             contract_name,
             proxy_name,
             custom,
+            override_entry_points,
         }
     }
 
@@ -696,32 +702,64 @@ impl<'a> MultitestHelpers<'a> {
 
     fn generate_impl_contract(&self) -> TokenStream {
         let Self {
-            contract, custom, ..
+            contract,
+            custom,
+            override_entry_points,
+            ..
         } = self;
         let sylvia = crate_module();
 
-        #[cfg(not(tarpaulin_include))]
-        let migrate_body = if self.is_migrate {
-            quote! {
-                #sylvia ::cw_std::from_slice::<MigrateMsg>(&msg)?
-                    .dispatch(self, (deps, env).into())
-                    .map_err(Into::into)
+        let instantiate_body = override_entry_points
+            .get_entry_point(MsgType::Instantiate)
+            .map(OverrideEntryPoint::emit_multitest_dispatch)
+            .unwrap_or_else(|| {
+                OverrideEntryPoint::emit_multitest_default_dispatch(MsgType::Instantiate)
+            });
+
+        let exec_body = override_entry_points
+            .get_entry_point(MsgType::Exec)
+            .map(OverrideEntryPoint::emit_multitest_dispatch)
+            .unwrap_or_else(|| OverrideEntryPoint::emit_multitest_default_dispatch(MsgType::Exec));
+
+        let query_body = override_entry_points
+            .get_entry_point(MsgType::Query)
+            .map(OverrideEntryPoint::emit_multitest_dispatch)
+            .unwrap_or_else(|| OverrideEntryPoint::emit_multitest_default_dispatch(MsgType::Query));
+
+        let sudo_body = override_entry_points
+            .get_entry_point(MsgType::Sudo)
+            .map(OverrideEntryPoint::emit_multitest_dispatch)
+            .unwrap_or_else(|| {
+                quote! {
+                    #sylvia ::anyhow::bail!("sudo not implemented for contract")
+                }
+            });
+
+        let migrate_body = match override_entry_points.get_entry_point(MsgType::Migrate) {
+            Some(entry_point) => entry_point.emit_multitest_dispatch(),
+            None if self.is_migrate => {
+                OverrideEntryPoint::emit_multitest_default_dispatch(MsgType::Migrate)
             }
-        } else {
-            quote! {
+            None => quote! {
                 #sylvia ::anyhow::bail!("migrate not implemented for contract")
-            }
+            },
         };
 
-        #[cfg(not(tarpaulin_include))]
-        let reply = if let Some(reply) = self.reply.as_ref() {
-            quote! {
-                self. #reply((deps, env).into(), msg).map_err(Into::into)
-            }
-        } else {
-            quote! {
-                #sylvia ::anyhow::bail!("reply not implemented for contract")
-            }
+        let reply_body = match override_entry_points.get_entry_point(MsgType::Reply) {
+            Some(entry_point) => entry_point.emit_multitest_dispatch(),
+            None => self
+                .reply
+                .as_ref()
+                .map(|reply| {
+                    quote! {
+                        self. #reply((deps, env).into(), msg).map_err(Into::into)
+                    }
+                })
+                .unwrap_or_else(|| {
+                    quote! {
+                        #sylvia ::anyhow::bail!("reply not implemented for contract")
+                    }
+                }),
         };
 
         let custom_msg = custom.msg();
@@ -737,9 +775,7 @@ impl<'a> MultitestHelpers<'a> {
                         info: #sylvia ::cw_std::MessageInfo,
                         msg: Vec<u8>,
                     ) -> #sylvia ::anyhow::Result<#sylvia ::cw_std::Response<#custom_msg>> {
-                        #sylvia ::cw_std::from_slice::<ContractExecMsg>(&msg)?
-                            .dispatch(self, (deps, env, info))
-                            .map_err(Into::into)
+                        #exec_body
                     }
 
                     fn instantiate(
@@ -749,9 +785,7 @@ impl<'a> MultitestHelpers<'a> {
                         info: #sylvia ::cw_std::MessageInfo,
                         msg: Vec<u8>,
                     ) -> #sylvia ::anyhow::Result<#sylvia ::cw_std::Response<#custom_msg>> {
-                        #sylvia ::cw_std::from_slice::<InstantiateMsg>(&msg)?
-                            .dispatch(self, (deps, env, info))
-                            .map_err(Into::into)
+                        #instantiate_body
                     }
 
                     fn query(
@@ -760,18 +794,16 @@ impl<'a> MultitestHelpers<'a> {
                         env: #sylvia ::cw_std::Env,
                         msg: Vec<u8>,
                     ) -> #sylvia ::anyhow::Result<#sylvia ::cw_std::Binary> {
-                        #sylvia ::cw_std::from_slice::<ContractQueryMsg>(&msg)?
-                            .dispatch(self, (deps, env))
-                            .map_err(Into::into)
+                        #query_body
                     }
 
                     fn sudo(
                         &self,
-                        _deps: #sylvia ::cw_std::DepsMut<#sylvia ::cw_std::Empty>,
-                        _env: #sylvia ::cw_std::Env,
-                        _msg: Vec<u8>,
+                        deps: #sylvia ::cw_std::DepsMut<#sylvia ::cw_std::Empty>,
+                        env: #sylvia ::cw_std::Env,
+                        msg: Vec<u8>,
                     ) -> #sylvia ::anyhow::Result<#sylvia ::cw_std::Response<#custom_msg>> {
-                        #sylvia ::anyhow::bail!("sudo not implemented for contract")
+                        #sudo_body
                     }
 
                     fn reply(
@@ -780,7 +812,7 @@ impl<'a> MultitestHelpers<'a> {
                         env: #sylvia ::cw_std::Env,
                         msg: #sylvia ::cw_std::Reply,
                     ) -> #sylvia ::anyhow::Result<#sylvia ::cw_std::Response<#custom_msg>> {
-                        #reply
+                        #reply_body
                     }
 
                     fn migrate(
diff --git a/sylvia/tests/entry_points.rs b/sylvia/tests/entry_points.rs
new file mode 100644
index 0000000..d7bbb2e
--- /dev/null
+++ b/sylvia/tests/entry_points.rs
@@ -0,0 +1,219 @@
+use contract::ContractExecMsg;
+use cosmwasm_schema::cw_serde;
+use cosmwasm_std::Coin;
+use cw_storage_plus::Item;
+
+#[cw_serde]
+pub struct CountResponse {
+    pub count: u32,
+}
+
+#[cw_serde]
+pub enum SudoMsg {
+    MoveFunds {
+        recipient: String,
+        amount: Vec<Coin>,
+    },
+}
+
+pub mod migrate {
+    use cosmwasm_schema::cw_serde;
+
+    #[cw_serde]
+    pub struct MigrateMsg {}
+}
+
+pub mod exec {
+    use cosmwasm_schema::cw_serde;
+    use cosmwasm_std::{DepsMut, Env, MessageInfo, Response, StdError, StdResult};
+    use sylvia::types::ExecCtx;
+
+    #[cw_serde]
+    pub enum UserExecMsg {
+        IncreaseByOne {},
+    }
+
+    pub fn increase_by_one(ctx: ExecCtx) -> StdResult<Response> {
+        crate::COUNTER.update(ctx.deps.storage, |count| -> Result<u32, StdError> {
+            Ok(count + 1)
+        })?;
+        Ok(Response::new())
+    }
+
+    #[cw_serde]
+    pub enum CustomExecMsg {
+        ContractExec(crate::ContractExecMsg),
+        CustomExec(UserExecMsg),
+    }
+
+    impl CustomExecMsg {
+        pub fn dispatch(self, ctx: (DepsMut, Env, MessageInfo)) -> StdResult<Response> {
+            match self {
+                CustomExecMsg::ContractExec(msg) => {
+                    msg.dispatch(&crate::contract::Contract::new(), ctx)
+                }
+                CustomExecMsg::CustomExec(_) => increase_by_one(ctx.into()),
+            }
+        }
+    }
+}
+
+const COUNTER: Item<u32> = Item::new("counter");
+
+pub mod entry_points {
+    use cosmwasm_std::{entry_point, DepsMut, Env, MessageInfo, Response, StdResult};
+
+    use crate::exec::CustomExecMsg;
+    use crate::migrate::MigrateMsg;
+    use crate::{SudoMsg, COUNTER};
+
+    #[entry_point]
+    pub fn sudo(deps: DepsMut, _env: Env, _msg: SudoMsg) -> StdResult<Response> {
+        COUNTER.save(deps.storage, &3)?;
+        Ok(Response::new())
+    }
+
+    #[entry_point]
+    pub fn migrate(deps: DepsMut, _env: Env, _msg: MigrateMsg) -> StdResult<Response> {
+        COUNTER.save(deps.storage, &5)?;
+        Ok(Response::new())
+    }
+
+    #[entry_point]
+    pub fn execute(
+        deps: DepsMut,
+        env: Env,
+        info: MessageInfo,
+        msg: CustomExecMsg,
+    ) -> StdResult<Response> {
+        msg.dispatch((deps, env, info))
+    }
+}
+
+mod contract {
+    use cosmwasm_std::{Response, StdError, StdResult};
+    use sylvia::contract;
+    use sylvia::types::{ExecCtx, InstantiateCtx, MigrateCtx, QueryCtx};
+
+    use crate::{CountResponse, COUNTER};
+
+    pub struct Contract {}
+
+    #[cfg(not(tarpaulin_include))]
+    #[contract]
+    #[sv::override_entry_point(sudo=crate::entry_points::sudo(crate::SudoMsg))]
+    #[sv::override_entry_point(migrate=crate::entry_points::migrate(crate::migrate::MigrateMsg))]
+    #[sv::override_entry_point(exec=crate::entry_points::execute(crate::exec::CustomExecMsg))]
+    #[allow(dead_code)]
+    impl Contract {
+        #[allow(clippy::new_without_default)]
+        #[allow(dead_code)]
+        pub fn new() -> Self {
+            Self {}
+        }
+
+        #[msg(instantiate)]
+        pub fn instantiate(&self, ctx: InstantiateCtx) -> StdResult<Response> {
+            COUNTER.save(ctx.deps.storage, &0)?;
+            Ok(Response::new())
+        }
+
+        #[msg(migrate)]
+        pub fn migrate(&self, _ctx: MigrateCtx) -> StdResult<Response> {
+            Ok(Response::new())
+        }
+
+        #[msg(query)]
+        pub fn count(&self, ctx: QueryCtx) -> StdResult<CountResponse> {
+            let count = COUNTER.load(ctx.deps.storage)?;
+            Ok(CountResponse { count })
+        }
+
+        #[msg(exec)]
+        pub fn increase_by_two(&self, ctx: ExecCtx) -> StdResult<Response> {
+            crate::COUNTER.update(ctx.deps.storage, |count| -> Result<u32, StdError> {
+                Ok(count + 2)
+            })?;
+            Ok(Response::new())
+        }
+    }
+}
+
+#[cfg(all(test, feature = "mt"))]
+mod tests {
+    use cosmwasm_std::Addr;
+    use cw_multi_test::Executor;
+    use sylvia::multitest::App;
+
+    use crate::contract::multitest_utils::CodeId;
+    use crate::contract::{ContractExecMsg, ExecMsg};
+    use crate::exec::{CustomExecMsg, UserExecMsg};
+    use crate::SudoMsg;
+
+    #[test]
+    fn overriden_entry_points_in_mt() {
+        let app = App::default();
+        let code_id = CodeId::store_code(&app);
+
+        let owner = "owner";
+
+        let contract = code_id
+            .instantiate()
+            .with_label("Contract")
+            .with_admin(Some(owner))
+            .call(owner)
+            .unwrap();
+
+        let count = contract.count().unwrap().count;
+        assert_eq!(count, 0);
+
+        let msg = SudoMsg::MoveFunds {
+            recipient: "recipient".to_string(),
+            amount: vec![],
+        };
+
+        contract
+            .app
+            .app_mut()
+            .wasm_sudo(contract.contract_addr.clone(), &msg)
+            .unwrap();
+
+        let count = contract.count().unwrap().count;
+        assert_eq!(count, 3);
+
+        contract.migrate().call(owner, code_id.code_id()).unwrap();
+        let count = contract.count().unwrap().count;
+        assert_eq!(count, 5);
+
+        // custom ExecMsg
+        let msg = CustomExecMsg::CustomExec(UserExecMsg::IncreaseByOne {});
+        (*contract.app)
+            .app_mut()
+            .execute_contract(
+                Addr::unchecked(owner),
+                contract.contract_addr.clone(),
+                &msg,
+                &[],
+            )
+            .unwrap();
+
+        let count = contract.count().unwrap().count;
+        assert_eq!(count, 6);
+
+        // custom ExecMsg
+        let msg =
+            CustomExecMsg::ContractExec(ContractExecMsg::Contract(ExecMsg::increase_by_two()));
+        (*contract.app)
+            .app_mut()
+            .execute_contract(
+                Addr::unchecked(owner),
+                contract.contract_addr.clone(),
+                &msg,
+                &[],
+            )
+            .unwrap();
+
+        let count = contract.count().unwrap().count;
+        assert_eq!(count, 8);
+    }
+}