diff --git a/README.md b/README.md
index df05e43..24f5b4b 100644
--- a/README.md
+++ b/README.md
@@ -221,9 +221,12 @@ $app->add(new Tuupola\Middleware\JwtAuthentication([
     "error" => function ($response, $arguments) {
         $data["status"] = "error";
         $data["message"] = $arguments["message"];
-        return $response
-            ->withHeader("Content-Type", "application/json")
-            ->getBody()->write(json_encode($data, JSON_UNESCAPED_SLASHES | JSON_PRETTY_PRINT));
+
+        $response->getBody()->write(
+            json_encode($data, JSON_UNESCAPED_SLASHES | JSON_PRETTY_PRINT)
+        );
+
+        return $response->withHeader("Content-Type", "application/json")
     }
 ]));
 ```
diff --git a/tests/JwtAuthenticationTest.php b/tests/JwtAuthenticationTest.php
index 05a6b44..4ac9652 100644
--- a/tests/JwtAuthenticationTest.php
+++ b/tests/JwtAuthenticationTest.php
@@ -751,8 +751,6 @@ public function testShouldCallAnonymousErrorFunction()
         $request = (new ServerRequestFactory)
             ->createServerRequest("GET", "https://example.com/api");
 
-        $dummy = null;
-
         $default = function (ServerRequestInterface $request) {
             $response = (new ResponseFactory)->createResponse();
             $response->getBody()->write("Success");
@@ -763,7 +761,9 @@ public function testShouldCallAnonymousErrorFunction()
             new JwtAuthentication([
                 "secret" => "supersecretkeyyoushouldnotcommit",
                 "error" => function (ResponseInterface $response, $arguments) use (&$dummy) {
-                    $dummy = true;
+                    $response->getBody()->write("error");
+                    return $response
+                        ->withHeader("X-Electrolytes", "Plants");
                 }
             ])
         ]);
@@ -771,8 +771,8 @@ public function testShouldCallAnonymousErrorFunction()
         $response = $collection->dispatch($request, $default);
 
         $this->assertEquals(401, $response->getStatusCode());
-        $this->assertEquals("", $response->getBody());
-        $this->assertTrue($dummy);
+        $this->assertEquals("Plants", $response->getHeaderLine("X-Electrolytes"));
+        $this->assertEquals("error", $response->getBody());
     }
 
     public function testShouldCallInvokableErrorClass()
@@ -798,6 +798,7 @@ public function testShouldCallInvokableErrorClass()
         $response = $collection->dispatch($request, $default);
 
         $this->assertEquals(402, $response->getStatusCode());
+        $this->assertEquals("Bar", $response->getHeaderLine("X-Foo"));
         $this->assertEquals(TestErrorHandler::class, $response->getBody());
     }
 
@@ -824,6 +825,7 @@ public function testShouldCallArrayNotationError()
         $response = $collection->dispatch($request, $default);
 
         $this->assertEquals(418, $response->getStatusCode());
+        $this->assertEquals("Foo", $response->getHeaderLine("X-Bar"));
         $this->assertEquals(TestErrorHandler::class, $response->getBody());
     }
 
diff --git a/tests/TestErrorHandler.php b/tests/TestErrorHandler.php
index 71f872a..e294a9f 100644
--- a/tests/TestErrorHandler.php
+++ b/tests/TestErrorHandler.php
@@ -41,7 +41,9 @@ public function __invoke(
         array $arguments
     ) {
         $response->getBody()->write(self::class);
-        return $response->withStatus(402);
+        return $response
+            ->withStatus(402)
+            ->withHeader("X-Foo", "Bar");
     }
 
     public static function error(
@@ -49,6 +51,8 @@ public static function error(
         array $arguments
     ) {
         $response->getBody()->write(self::class);
-        return $response->withStatus(418);
+        return $response
+            ->withStatus(418)
+            ->withHeader("X-Bar", "Foo");
     }
 }