From 2de0ac59e6e6e22d9dd193ebca007906881e60e7 Mon Sep 17 00:00:00 2001
From: Nelson Vides <nelson.vides@erlang-solutions.com>
Date: Mon, 12 Aug 2024 11:53:20 +0200
Subject: [PATCH] Correctly implement the code_change callback

---
 src/wpool_process.erl        | 5 +++--
 test/wpool_process_SUITE.erl | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/wpool_process.erl b/src/wpool_process.erl
index 63c45e3..9105e6e 100644
--- a/src/wpool_process.erl
+++ b/src/wpool_process.erl
@@ -170,14 +170,15 @@ terminate(Reason, State) ->
     end.
 
 %% @private
--spec code_change(string(), state(), any()) -> {ok, state()} | {error, term()}.
+-spec code_change(string() | {down, string()}, state(), any()) ->
+                     {ok, state()} | {error, term()}.
 code_change(OldVsn, #state{mod = #callback_cache{module = Mod}} = State, Extra) ->
     case erlang:function_exported(Mod, code_change, 3) of
         true ->
             case Mod:code_change(OldVsn, State#state.state, Extra) of
                 {ok, NewState} ->
                     {ok, State#state{state = NewState}};
-                Error ->
+                {error, Error} ->
                     {error, Error}
             end;
         _ ->
diff --git a/test/wpool_process_SUITE.erl b/test/wpool_process_SUITE.erl
index 538284c..1e63f8b 100644
--- a/test/wpool_process_SUITE.erl
+++ b/test/wpool_process_SUITE.erl
@@ -334,7 +334,7 @@ complete_coverage(_Config) ->
     ct:comment("Code Change"),
     {ok, State} = wpool_process:init({complete_coverage, echo_server, {ok, state}, []}),
     {ok, _} = wpool_process:code_change("oldvsn", State, {ok, state}),
-    {error, bad} = wpool_process:code_change("oldvsn", State, bad),
+    {error, bad} = wpool_process:code_change("oldvsn", State, {error, bad}),
 
     {comment, []}.