From 3a8395a432cefca0fca7d77826583dad6d9711e1 Mon Sep 17 00:00:00 2001
From: michaelboulton <michaelboulton@users.noreply.github.com>
Date: Sat, 10 Feb 2024 18:27:52 +0000
Subject: [PATCH] Fix saving in grpc (#916)

---
 example/grpc/test_grpc.tavern.yaml | 37 ++++++++++++++++++++++++++++++
 tavern/_plugins/grpc/response.py   |  2 +-
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/example/grpc/test_grpc.tavern.yaml b/example/grpc/test_grpc.tavern.yaml
index 08382346..d9c9883d 100644
--- a/example/grpc/test_grpc.tavern.yaml
+++ b/example/grpc/test_grpc.tavern.yaml
@@ -26,6 +26,43 @@ stages:
 
 ---
 
+test_name: Test grpc saving
+
+includes:
+  - !include common.yaml
+
+grpc:
+  connect:
+    <<: *grpc_connect
+  proto:
+    module: helloworld_v1_precompiled_pb2_grpc
+
+stages:
+  - name: Echo text
+    grpc_request:
+      service: helloworld.v1.Greeter/SayHello
+      body:
+        name: "John"
+    grpc_response:
+      status: "OK"
+      body:
+        message: "Hello, John!"
+      save:
+        body:
+          received_message: message
+
+  - name: Echo text
+    grpc_request:
+      service: helloworld.v1.Greeter/SayHello
+      body:
+        name: "{received_message}"
+    grpc_response:
+      status: "OK"
+      body:
+        message: "Hello, Hello, John!!"
+
+---
+
 test_name: Test trying to connect using an invalid option
 
 includes:
diff --git a/tavern/_plugins/grpc/response.py b/tavern/_plugins/grpc/response.py
index b9c8849d..f2e9c000 100644
--- a/tavern/_plugins/grpc/response.py
+++ b/tavern/_plugins/grpc/response.py
@@ -49,7 +49,7 @@ def __init__(
         expected: Union[_GRPCExpected, Mapping],
         test_block_config: TestConfig,
     ) -> None:
-        check_expected_keys({"body", "status", "details"}, expected)
+        check_expected_keys({"body", "status", "details", "save"}, expected)
         super().__init__(name, expected, test_block_config)
 
         self._client = client