From 1ddd078d324fde76ca178a97321e22e569b5ff4f Mon Sep 17 00:00:00 2001
From: Pawan Dhananjay <pawandhananjay@gmail.com>
Date: Tue, 7 May 2024 16:13:36 -0700
Subject: [PATCH] Handle new engine api methods in mock EL

---
 .../src/test_utils/handle_rpc.rs              | 34 ++++++++++---------
 1 file changed, 18 insertions(+), 16 deletions(-)

diff --git a/beacon_node/execution_layer/src/test_utils/handle_rpc.rs b/beacon_node/execution_layer/src/test_utils/handle_rpc.rs
index 1dc8f0ab83..5e31646d31 100644
--- a/beacon_node/execution_layer/src/test_utils/handle_rpc.rs
+++ b/beacon_node/execution_layer/src/test_utils/handle_rpc.rs
@@ -111,20 +111,11 @@ pub async fn handle_rpc<E: EthSpec>(
                             .map(|jep| JsonExecutionPayload::V1(jep))
                     })
                     .map_err(|s| (s, BAD_PARAMS_ERROR_CODE))?,
-                ENGINE_NEW_PAYLOAD_V3 => get_param::<JsonExecutionPayloadV4<E>>(params, 0)
+                ENGINE_NEW_PAYLOAD_V3 => get_param::<JsonExecutionPayloadV3<E>>(params, 0)
+                    .map(|jep| JsonExecutionPayload::V3(jep))
+                    .map_err(|s| (s, BAD_PARAMS_ERROR_CODE))?,
+                ENGINE_NEW_PAYLOAD_V4 => get_param::<JsonExecutionPayloadV4<E>>(params, 0)
                     .map(|jep| JsonExecutionPayload::V4(jep))
-                    .or_else(|_| {
-                        get_param::<JsonExecutionPayloadV3<E>>(params, 0)
-                            .map(|jep| JsonExecutionPayload::V3(jep))
-                            .or_else(|_| {
-                                get_param::<JsonExecutionPayloadV2<E>>(params, 0)
-                                    .map(|jep| JsonExecutionPayload::V2(jep))
-                                    .or_else(|_| {
-                                        get_param::<JsonExecutionPayloadV1<E>>(params, 0)
-                                            .map(|jep| JsonExecutionPayload::V1(jep))
-                                    })
-                            })
-                    })
                     .map_err(|s| (s, BAD_PARAMS_ERROR_CODE))?,
                 _ => unreachable!(),
             };
@@ -190,7 +181,10 @@ pub async fn handle_rpc<E: EthSpec>(
                     }
                 }
                 ForkName::Electra => {
-                    if method == ENGINE_NEW_PAYLOAD_V1 || method == ENGINE_NEW_PAYLOAD_V2 {
+                    if method == ENGINE_NEW_PAYLOAD_V1
+                        || method == ENGINE_NEW_PAYLOAD_V2
+                        || method == ENGINE_NEW_PAYLOAD_V3
+                    {
                         return Err((
                             format!("{} called after Electra fork!", method),
                             GENERIC_ERROR_CODE,
@@ -259,7 +253,10 @@ pub async fn handle_rpc<E: EthSpec>(
 
             Ok(serde_json::to_value(JsonPayloadStatusV1::from(response)).unwrap())
         }
-        ENGINE_GET_PAYLOAD_V1 | ENGINE_GET_PAYLOAD_V2 | ENGINE_GET_PAYLOAD_V3 => {
+        ENGINE_GET_PAYLOAD_V1
+        | ENGINE_GET_PAYLOAD_V2
+        | ENGINE_GET_PAYLOAD_V3
+        | ENGINE_GET_PAYLOAD_V4 => {
             let request: JsonPayloadIdRequest =
                 get_param(params, 0).map_err(|s| (s, BAD_PARAMS_ERROR_CODE))?;
             let id = request.into();
@@ -309,7 +306,9 @@ pub async fn handle_rpc<E: EthSpec>(
                 .read()
                 .get_fork_at_timestamp(response.timestamp())
                 == ForkName::Electra
-                && method == ENGINE_GET_PAYLOAD_V1
+                && (method == ENGINE_GET_PAYLOAD_V1
+                    || method == ENGINE_GET_PAYLOAD_V2
+                    || method == ENGINE_GET_PAYLOAD_V3)
             {
                 return Err((
                     format!("{} called after Electra fork!", method),
@@ -353,6 +352,9 @@ pub async fn handle_rpc<E: EthSpec>(
                         })
                         .unwrap()
                     }
+                    _ => unreachable!(),
+                }),
+                ENGINE_GET_PAYLOAD_V4 => Ok(match JsonExecutionPayload::from(response) {
                     JsonExecutionPayload::V4(execution_payload) => {
                         serde_json::to_value(JsonGetPayloadResponseV4 {
                             execution_payload,