@layerzerolabs/framework-solana-anchor-trait 0.2.62 → 0.2.64

package/tests/generation.rs

@@ -3,7 +3,7 @@ mod common;
 #[cfg(test)]
 mod tests {
     use super::common::*;
-    use quote::ToTokens;
+    use quote::{quote, ToTokens};
 
     // -----------------------------------------------------------------------
     // Tests — override level resolution
@@ -11,24 +11,21 @@ mod tests {
 
     #[test]
     fn default_level_emits_context_and_entrypoint() {
-        // All defaults — no overrides except required lz_receive
         let module: syn::ItemMod = syn::parse_quote! {
             mod my_program {
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
             }
         };
         let gen = generate_ok(&module);
 
-        // Default impls should be generated for each provided instruction
         let structs = parse_structs(&gen.default_impls);
         assert!(find_struct(&structs, "StdSetPeer").is_some(), "should emit StdSetPeer");
         assert!(find_struct(&structs, "StdGetPeer").is_some(), "should emit StdGetPeer");
         assert!(find_struct(&structs, "StdToken").is_some(), "should emit StdToken");
 
-        // Entrypoints should include all instructions
         let fns = parse_fns(&gen.entrypoints);
         assert!(find_fn(&fns, "set_peer").is_some(), "should have set_peer entrypoint");
         assert!(find_fn(&fns, "get_peer").is_some(), "should have get_peer entrypoint");
@@ -41,7 +38,7 @@ mod tests {
         let module: syn::ItemMod = syn::parse_quote! {
             mod my_program {
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
             }
@@ -55,11 +52,11 @@ mod tests {
         assert!(is_mut, "set_peer should have `mut ctx`");
         assert_eq!(ctx_type, "StdSetPeer", "should use default context type");
 
-        // Default handler called by reference: StdSetPeer::apply(&mut ctx, ...)
+        // Default handler called by reference: StdSetPeer::apply(&mut ctx, &params)
         let body = set_peer.block.to_token_stream().to_string();
         assert!(
-            body.contains("StdSetPeer :: apply (& mut ctx"),
-            "mutable default should pass &mut ctx to handler: {body}"
+            body.contains("StdSetPeer :: apply (& mut ctx , & params)"),
+            "mutable default should pass &mut ctx and &params to handler: {body}"
         );
     }
 
@@ -68,7 +65,7 @@ mod tests {
         let module: syn::ItemMod = syn::parse_quote! {
             mod my_program {
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
             }
@@ -82,40 +79,38 @@ mod tests {
         assert!(!is_mut, "view default should use `ctx` (no mut)");
         assert_eq!(ctx_type, "StdGetPeer", "should use default context type");
 
-        // Default handler called by immutable reference: StdGetPeer::apply(&ctx, ...)
+        // Default handler called by immutable reference: StdGetPeer::apply(&ctx, &params)
         let body = get_peer.block.to_token_stream().to_string();
         assert!(
-            body.contains("StdGetPeer :: apply (& ctx"),
-            "view default should pass &ctx to handler: {body}"
+            body.contains("StdGetPeer :: apply (& ctx , & params)"),
+            "view default should pass &ctx and &params to handler: {body}"
         );
     }
 
     #[test]
-    fn handler_override_emits_default_context() {
-        // Override set_peer but reuse StdSetPeer context
+    fn override_with_default_context_emits_default() {
+        // Override set_peer but reuse StdSetPeer context.
         let module: syn::ItemMod = syn::parse_quote! {
             mod my_program {
                 #[test_instruction]
-                pub fn set_peer(ctx: Context<StdSetPeer>, params: SetPeerParams) -> Result<()> {
+                pub fn set_peer(ctx: &mut Context<StdSetPeer>, params: &SetPeerParams) -> Result<()> {
                     Ok(())
                 }
 
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
             }
         };
         let gen = generate_ok(&module);
 
-        // Default impl should still be generated (handler reuses it)
         let structs = parse_structs(&gen.default_impls);
         assert!(
             find_struct(&structs, "StdSetPeer").is_some(),
-            "handler override should still emit StdSetPeer context"
+            "override should still emit StdSetPeer context"
         );
 
-        // Entrypoint should dispatch to wrapper
         let fns = parse_fns(&gen.entrypoints);
         let set_peer = find_fn(&fns, "set_peer").expect("set_peer entrypoint should exist");
         let body = set_peer.block.to_token_stream().to_string();
@@ -126,29 +121,29 @@ mod tests {
     }
 
     #[test]
-    fn full_override_skips_default_context() {
-        // Override set_peer with a custom context type (not StdSetPeer)
+    fn override_with_custom_context_still_emits_default() {
+        // Override set_peer with a custom context type (not StdSetPeer).
+        // Default struct is always emitted for provided instructions, even when
+        // the override handler uses a different context type.
         let module: syn::ItemMod = syn::parse_quote! {
             mod my_program {
                 #[test_instruction]
-                pub fn set_peer(ctx: Context<MyCustomAccounts>, params: SetPeerParams) -> Result<()> {
+                pub fn set_peer(ctx: &mut Context<MyCustomAccounts>, params: &SetPeerParams) -> Result<()> {
                     Ok(())
                 }
 
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
             }
         };
         let gen = generate_ok(&module);
 
-        // Full override should NOT emit default impl for set_peer.
-        // Other provided instructions (get_peer, token) still get their defaults.
         let structs = parse_structs(&gen.default_impls);
         assert!(
-            find_struct(&structs, "StdSetPeer").is_none(),
-            "full override should skip default impl for set_peer"
+            find_struct(&structs, "StdSetPeer").is_some(),
+            "override should still emit StdSetPeer (always available)"
         );
         assert!(
             find_struct(&structs, "StdGetPeer").is_some(),
@@ -156,14 +151,12 @@ mod tests {
         );
         assert!(find_struct(&structs, "StdToken").is_some(), "token should still have default");
 
-        // Entrypoint should use the custom context type
         let fns = parse_fns(&gen.entrypoints);
         let set_peer = find_fn(&fns, "set_peer").expect("set_peer entrypoint should exist");
         let (is_mut, ctx_type) = first_param_context_type(set_peer);
         assert!(is_mut, "set_peer should have `mut ctx`");
-        assert_eq!(ctx_type, "MyCustomAccounts", "full override should use custom context type");
+        assert_eq!(ctx_type, "MyCustomAccounts", "override should use custom context type");
 
-        // Should dispatch to wrapper
         let body = set_peer.block.to_token_stream().to_string();
         assert!(
             body.contains("__test_instruction__set_peer_handler :: set_peer"),
@@ -172,16 +165,16 @@ mod tests {
     }
 
     #[test]
-    fn override_receives_params_by_value() {
+    fn override_receives_params_by_ref() {
         let module: syn::ItemMod = syn::parse_quote! {
             mod my_program {
                 #[test_instruction]
-                pub fn set_peer(ctx: Context<MyAccounts>, params: SetPeerParams) -> Result<()> {
+                pub fn set_peer(ctx: &mut Context<MyAccounts>, params: &SetPeerParams) -> Result<()> {
                     Ok(())
                 }
 
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
             }
@@ -190,11 +183,10 @@ mod tests {
         let fns = parse_fns(&gen.entrypoints);
         let set_peer = find_fn(&fns, "set_peer").expect("set_peer entrypoint should exist");
 
-        // Override handlers get (ctx, params) by value — no & or &mut on the call
         let body = set_peer.block.to_token_stream().to_string();
         assert!(
-            body.contains("set_peer (ctx , params)"),
-            "override should pass ctx and params by value: {body}"
+            body.contains("set_peer (& mut ctx , & params)"),
+            "override should pass ctx by &mut and params by &: {body}"
         );
     }
 
@@ -203,7 +195,7 @@ mod tests {
         let module: syn::ItemMod = syn::parse_quote! {
             mod my_program {
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
             }
@@ -212,7 +204,6 @@ mod tests {
         let fns = parse_fns(&gen.entrypoints);
         let token = find_fn(&fns, "token").expect("token entrypoint should exist");
 
-        // Token instruction has no params — entrypoint should only have ctx
         assert_eq!(token.sig.inputs.len(), 1, "no-params instruction should have only ctx");
         let (is_mut, ctx_type) = first_param_context_type(token);
         assert!(!is_mut, "token is a view — should not be mut");
@@ -224,18 +215,18 @@ mod tests {
     // -----------------------------------------------------------------------
 
     #[test]
-    fn view_override_with_params_gets_immutable_ctx() {
-        // Override get_peer (view, has params) — entrypoint should have
-        // immutable ctx and dispatch to wrapper with (ctx, params) by value.
+    fn view_override_with_params_gets_immutable_ref() {
+        // Override get_peer (view, has params) — entrypoint must have immutable ctx
+        // and dispatch to wrapper with (&ctx, &params).
         let module: syn::ItemMod = syn::parse_quote! {
             mod my_program {
                 #[test_instruction]
-                pub fn get_peer(ctx: Context<MyGetPeerAccounts>, params: GetPeerParams) -> Result<[u8; 32]> {
+                pub fn get_peer(ctx: &Context<MyGetPeerAccounts>, params: &GetPeerParams) -> Result<[u8; 32]> {
                     Ok([0u8; 32])
                 }
 
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
             }
@@ -244,32 +235,29 @@ mod tests {
         let fns = parse_fns(&gen.entrypoints);
         let get_peer = find_fn(&fns, "get_peer").expect("get_peer entrypoint should exist");
 
-        // View override: ctx should NOT be mut
         let (is_mut, ctx_type) = first_param_context_type(get_peer);
         assert!(!is_mut, "view override should have immutable ctx");
         assert_eq!(ctx_type, "MyGetPeerAccounts", "should use override's context type");
 
-        // Dispatches to wrapper by value
         let body = get_peer.block.to_token_stream().to_string();
         assert!(
-            body.contains("get_peer (ctx , params)"),
-            "view override should pass ctx and params by value: {body}"
+            body.contains("get_peer (& ctx , & params)"),
+            "view override should pass ctx by & and params by &: {body}"
         );
     }
 
     #[test]
-    fn override_no_params_receives_ctx_by_value() {
-        // Override the token instruction (no params) — entrypoint should
-        // dispatch to wrapper with just `ctx` by value.
+    fn override_no_params_receives_ctx_by_ref() {
+        // Override token (no params, view) — entrypoint must dispatch with just `&ctx`.
         let module: syn::ItemMod = syn::parse_quote! {
             mod my_program {
                 #[test_instruction]
-                pub fn token(ctx: Context<MyTokenAccounts>) -> Result<Pubkey> {
+                pub fn token(ctx: &Context<MyTokenAccounts>) -> Result<Pubkey> {
                     Ok(Pubkey::default())
                 }
 
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
             }
@@ -278,13 +266,9 @@ mod tests {
         let fns = parse_fns(&gen.entrypoints);
         let token = find_fn(&fns, "token").expect("token entrypoint should exist");
 
-        // No-params override: entrypoint has only ctx, dispatches wrapper::handler(ctx)
         assert_eq!(token.sig.inputs.len(), 1, "no-params override should have only ctx");
         let body = token.block.to_token_stream().to_string();
-        assert!(
-            body.contains("token (ctx)"),
-            "no-params override should pass ctx by value: {body}"
-        );
+        assert!(body.contains("token (& ctx)"), "no-params override should pass ctx by &: {body}");
     }
 
     #[test]
@@ -292,7 +276,7 @@ mod tests {
         let module: syn::ItemMod = syn::parse_quote! {
             mod my_program {
                 #[test_instruction]
-                pub fn lz_receive(ctx: Context<MyReceive>, params: LzReceiveParams) -> Result<()> {
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
                     Ok(())
                 }
             }
@@ -300,17 +284,15 @@ mod tests {
         let gen = generate_ok(&module);
         let fns = parse_fns(&gen.entrypoints);
 
-        // get_peer returns [u8; 32] per spec
+        // Return types come from the spec, not the override handler.
         let get_peer = find_fn(&fns, "get_peer").expect("get_peer should exist");
         let ret = get_peer.sig.output.to_token_stream().to_string();
         assert!(ret.contains("[u8 ; 32]"), "get_peer return type should be [u8; 32]: {ret}");
 
-        // token returns Pubkey per spec
         let token = find_fn(&fns, "token").expect("token should exist");
         let ret = token.sig.output.to_token_stream().to_string();
         assert!(ret.contains("Pubkey"), "token return type should be Pubkey: {ret}");
 
-        // set_peer returns () per spec
         let set_peer = find_fn(&fns, "set_peer").expect("set_peer should exist");
         let ret = set_peer.sig.output.to_token_stream().to_string();
         assert!(ret.contains("()"), "set_peer return type should be (): {ret}");
@@ -323,8 +305,8 @@ mod tests {
             mod my_program {
                 #[test_instruction]
                 pub fn lz_receive<'info>(
-                    ctx: Context<'_, '_, '_, 'info, MyReceive<'info>>,
-                    params: LzReceiveParams,
+                    ctx: &mut Context<'_, '_, '_, 'info, MyReceive<'info>>,
+                    params: &LzReceiveParams,
                 ) -> Result<()> {
                     Ok(())
                 }
@@ -341,13 +323,103 @@ mod tests {
         );
     }
 
+    // -----------------------------------------------------------------------
+    // Tests — extend rules
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn extend_rules_generated_for_provided_instructions() {
+        let module: syn::ItemMod = syn::parse_quote! {
+            mod my_program {
+                #[test_instruction]
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
+                    Ok(())
+                }
+            }
+        };
+        let gen = generate_ok(&module);
+
+        let rules = gen.extend_rules.to_string();
+        assert!(
+            !rules.is_empty(),
+            "extend_rules should be non-empty when provided instructions exist"
+        );
+        assert!(rules.contains("set_peer"), "should contain a rule for set_peer");
+        assert!(rules.contains("get_peer"), "should contain a rule for get_peer");
+        assert!(rules.contains("token"), "should contain a rule for token");
+        assert!(
+            !rules.contains("lz_receive"),
+            "required instruction should NOT have an extend rule"
+        );
+    }
+
+    #[test]
+    fn extend_rule_contains_base_fields_and_extra_placeholder() {
+        let module: syn::ItemMod = syn::parse_quote! {
+            mod my_program {
+                #[test_instruction]
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
+                    Ok(())
+                }
+            }
+        };
+        let gen = generate_ok(&module);
+
+        let rules = gen.extend_rules.to_string();
+        assert!(rules.contains("placeholder"), "extend rules should include base fields");
+        assert!(
+            rules.contains("extra"),
+            "extend rules should include $($extra)* for appended accounts"
+        );
+    }
+
+    #[test]
+    fn extend_rule_preserves_struct_attributes() {
+        let module: syn::ItemMod = syn::parse_quote! {
+            mod my_program {
+                #[test_instruction]
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
+                    Ok(())
+                }
+            }
+        };
+        let gen = generate_ok(&module);
+
+        let rules = gen.extend_rules.to_string();
+        // Compare against the tokenized form rather than a raw string to be precise.
+        let expected = quote! { #[derive(Accounts)] }.to_string();
+        assert!(
+            rules.contains(&expected),
+            "extend rules should preserve struct attributes `{expected}`: {rules}"
+        );
+    }
+
+    #[test]
+    fn extend_rule_handler_uses_context_self() {
+        let module: syn::ItemMod = syn::parse_quote! {
+            mod my_program {
+                #[test_instruction]
+                pub fn lz_receive(ctx: &mut Context<MyReceive>, params: &LzReceiveParams) -> Result<()> {
+                    Ok(())
+                }
+            }
+        };
+        let gen = generate_ok(&module);
+
+        let rules = gen.extend_rules.to_string();
+        assert!(
+            rules.contains("Context < Self >"),
+            "extend rule handlers must use Context<Self> so they work for any extended struct name: {rules}"
+        );
+    }
+
     // -----------------------------------------------------------------------
     // Tests — signature validation errors
     // -----------------------------------------------------------------------
 
     #[test]
     fn error_required_instruction_missing_handler() {
-        // lz_receive is required (no default) but not overridden
+        // lz_receive is required (no default) but not overridden.
         let module: syn::ItemMod = syn::parse_quote! {
             mod my_program {}
         };