From 4914a075ea3a2822562df917225cad4720ff565e Mon Sep 17 00:00:00 2001
From: Joel Smith <joelsmith.2019@gmail.com>
Date: Wed, 11 Sep 2024 12:17:51 -0500
Subject: [PATCH] Add PFM Format Unit Tests

---
 .../types/ibc_hooks_test.go                   |   1 +
 .../x/ibc-composability-mw/types/pfm_test.go  | 111 +++++++++++++++++-
 2 files changed, 111 insertions(+), 1 deletion(-)

diff --git a/wormchain/x/ibc-composability-mw/types/ibc_hooks_test.go b/wormchain/x/ibc-composability-mw/types/ibc_hooks_test.go
index e61de326b6..c63a0a11b5 100644
--- a/wormchain/x/ibc-composability-mw/types/ibc_hooks_test.go
+++ b/wormchain/x/ibc-composability-mw/types/ibc_hooks_test.go
@@ -68,6 +68,7 @@ func TestFormatIbcHooksMemo(t *testing.T) {
 
 		if tc.shouldErr {
 			require.Error(t, err)
+			continue
 		} else {
 			require.NoError(t, err)
 			require.NotNil(t, res)
diff --git a/wormchain/x/ibc-composability-mw/types/pfm_test.go b/wormchain/x/ibc-composability-mw/types/pfm_test.go
index d9be07a2b5..cf30f9399c 100644
--- a/wormchain/x/ibc-composability-mw/types/pfm_test.go
+++ b/wormchain/x/ibc-composability-mw/types/pfm_test.go
@@ -1,7 +1,116 @@
 package types_test
 
-import "testing"
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"github.com/wormhole-foundation/wormchain/x/ibc-composability-mw/types"
+)
 
 func TestFormatPfmMemo(t *testing.T) {
+	for _, tc := range []struct {
+		payload   types.ParsedPayload
+		queryResp types.IbcTranslatorQueryRsp
+		timeout   time.Duration
+		retries   uint8
+		shouldErr bool
+	}{
+		// Normal w/ no payload
+		{
+			payload: types.ParsedPayload{
+				NoPayload: true,
+				ChainId:   1,
+				Recipient: []byte{'a', 'b', 'c'},
+				Fee:       "0uworm",
+				Nonce:     1,
+				Payload:   nil,
+			},
+			queryResp: types.IbcTranslatorQueryRsp{
+				Channel: "channel",
+			},
+			timeout:   time.Hour,
+			retries:   3,
+			shouldErr: false,
+		},
+		// Provide payload when unnecessary
+		{
+			payload: types.ParsedPayload{
+				NoPayload: true,
+				ChainId:   1,
+				Recipient: []byte{'a', 'b', 'c'},
+				Fee:       "0uworm",
+				Nonce:     1,
+				Payload:   []byte("{\"payload\":\"data\"}"),
+			},
+			queryResp: types.IbcTranslatorQueryRsp{
+				Channel: "channel",
+			},
+			timeout:   time.Hour,
+			retries:   3,
+			shouldErr: false,
+		},
+		// Normal w/ payload
+		{
+			payload: types.ParsedPayload{
+				NoPayload: false,
+				ChainId:   1,
+				Recipient: []byte{'a', 'b', 'c'},
+				Fee:       "0uworm",
+				Nonce:     1,
+				Payload:   []byte("{\"payload\":\"data\"}"),
+			},
+			queryResp: types.IbcTranslatorQueryRsp{
+				Channel: "channel-34",
+			},
+			timeout:   time.Minute,
+			retries:   21,
+			shouldErr: false,
+		},
+		// Nil payload should not err
+		{
+			payload: types.ParsedPayload{
+				NoPayload: true,
+				ChainId:   1,
+				Recipient: []byte{'a', 'b', 'c'},
+				Fee:       "0uworm",
+				Nonce:     1,
+				Payload:   nil,
+			},
+			queryResp: types.IbcTranslatorQueryRsp{
+				Channel: "channel",
+			},
+			timeout:   time.Microsecond,
+			retries:   255,
+			shouldErr: false,
+		},
+	} {
+
+		// turn the query response into bytes
+		queryRespBz, err := json.Marshal(tc.queryResp)
+		require.NoError(t, err)
+
+		res, err := types.FormatPfmMemo(tc.payload, queryRespBz, tc.timeout, tc.retries)
+
+		if tc.shouldErr {
+			require.Error(t, err)
+			continue
+		} else {
+			require.NoError(t, err)
+			require.NotNil(t, res)
+		}
+
+		// convert response back to packet metadata
+		var packetMetadata types.PacketMetadata
+		err = json.Unmarshal([]byte(res), &packetMetadata)
+		require.NoError(t, err)
 
+		// validation checks
+		require.Equal(t, string(tc.payload.Recipient), packetMetadata.Forward.Receiver)
+		require.Equal(t, "transfer", packetMetadata.Forward.Port)
+		require.Equal(t, tc.queryResp.Channel, packetMetadata.Forward.Channel)
+		require.Equal(t, tc.timeout, packetMetadata.Forward.Timeout)
+		require.Equal(t, &tc.retries, packetMetadata.Forward.Retries)
+	}
 }