Dedupe and add unit test

solana-labs · Sep 3, 2021 · f2cbacc · f2cbacc
1 parent a579d18
commit f2cbacc
Showing 1 changed file with 102 additions and 12 deletions.
diff --git a/transaction-status/src/extract_memos.rs b/transaction-status/src/extract_memos.rs
@@ -27,6 +27,14 @@ pub fn extract_and_fmt_memos<T: ExtractMemos>(message: &T) -> Option<String> {
     }
 }
 
+fn maybe_push_parsed_memo(memos: &mut Vec<String>, program_id: Pubkey, data: &[u8]) {
+    if program_id == spl_memo_id_v1() || program_id == spl_memo_id_v3() {
+        let memo_len = data.len();
+        let parsed_memo = parse_memo_data(data).unwrap_or_else(|_| "(unparseable)".to_string());
+        memos.push(format!("[{}] {}", memo_len, parsed_memo));
+    }
+}
+
 pub trait ExtractMemos {
     fn extract_memos(&self) -> Vec<String>;
 }
@@ -39,12 +47,7 @@ impl ExtractMemos for Message {
         {
             for instruction in &self.instructions {
                 let program_id = self.account_keys[instruction.program_id_index as usize];
-                if program_id == spl_memo_id_v1() || program_id == spl_memo_id_v3() {
-                    let memo_len = instruction.data.len();
-                    let parsed_memo = parse_memo_data(&instruction.data)
-                        .unwrap_or_else(|_| "(unparseable)".to_string());
-                    memos.push(format!("[{}] {}", memo_len, parsed_memo));
-                }
+                maybe_push_parsed_memo(&mut memos, program_id, &instruction.data);
             }
         }
         memos
@@ -59,14 +62,101 @@ impl ExtractMemos for SanitizedMessage {
             .any(|&pubkey| pubkey == spl_memo_id_v1() || pubkey == spl_memo_id_v3())
         {
             for (program_id, instruction) in self.program_instructions_iter() {
-                if *program_id == spl_memo_id_v1() || *program_id == spl_memo_id_v3() {
-                    let memo_len = instruction.data.len();
-                    let parsed_memo = parse_memo_data(&instruction.data)
-                        .unwrap_or_else(|_| "(unparseable)".to_string());
-                    memos.push(format!("[{}] {}", memo_len, parsed_memo));
-                }
+                maybe_push_parsed_memo(&mut memos, *program_id, &instruction.data);
             }
         }
         memos
     }
 }
+
+#[cfg(test)]
+mod test {
+    use {
+        super::*,
+        solana_sdk::{
+            hash::Hash,
+            instruction::CompiledInstruction,
+            message::{v0, MappedAddresses, MappedMessage, MessageHeader},
+        },
+    };
+
+    #[test]
+    fn test_extract_memos() {
+        let fee_payer = Pubkey::new_unique();
+        let another_program_id = Pubkey::new_unique();
+        let memo0 = "Test memo";
+        let memo1 = "🦖";
+        let expected_memos = vec![
+            format!("[{}] {}", memo0.len(), memo0),
+            format!("[{}] {}", memo1.len(), memo1),
+        ];
+        let message = Message::new_with_compiled_instructions(
+            1,
+            0,
+            3,
+            vec![
+                fee_payer,
+                spl_memo_id_v1(),
+                another_program_id,
+                spl_memo_id_v3(),
+            ],
+            Hash::default(),
+            vec![
+                CompiledInstruction {
+                    program_id_index: 1,
+                    accounts: vec![],
+                    data: memo0.as_bytes().to_vec(),
+                },
+                CompiledInstruction {
+                    program_id_index: 2,
+                    accounts: vec![],
+                    data: memo1.as_bytes().to_vec(),
+                },
+                CompiledInstruction {
+                    program_id_index: 3,
+                    accounts: vec![],
+                    data: memo1.as_bytes().to_vec(),
+                },
+            ],
+        );
+        assert_eq!(message.extract_memos(), expected_memos);
+
+        let sanitized_message = SanitizedMessage::Legacy(message);
+        assert_eq!(sanitized_message.extract_memos(), expected_memos);
+
+        let mapped_message = MappedMessage {
+            message: v0::Message {
+                header: MessageHeader {
+                    num_required_signatures: 1,
+                    num_readonly_signed_accounts: 0,
+                    num_readonly_unsigned_accounts: 3,
+                },
+                account_keys: vec![fee_payer],
+                instructions: vec![
+                    CompiledInstruction {
+                        program_id_index: 1,
+                        accounts: vec![],
+                        data: memo0.as_bytes().to_vec(),
+                    },
+                    CompiledInstruction {
+                        program_id_index: 2,
+                        accounts: vec![],
+                        data: memo1.as_bytes().to_vec(),
+                    },
+                    CompiledInstruction {
+                        program_id_index: 3,
+                        accounts: vec![],
+                        data: memo1.as_bytes().to_vec(),
+                    },
+                ],
+                ..v0::Message::default()
+            },
+            mapped_addresses: MappedAddresses {
+                writable: vec![],
+                readonly: vec![spl_memo_id_v1(), another_program_id, spl_memo_id_v3()],
+            },
+        };
+        let sanitized_mapped_message = SanitizedMessage::V0(mapped_message);
+        assert_eq!(sanitized_mapped_message.extract_memos(), expected_memos);
+    }
+}