From 72b8a928dc57a557c01c25f3b2a70cd483990b4b Mon Sep 17 00:00:00 2001
From: Tim Holm <tim.holm@nitric.io>
Date: Tue, 25 Oct 2022 08:56:21 +1100
Subject: [PATCH] test: Update storage tests.

---
 nitric/api/storage.py     |  4 ++--
 tests/api/test_storage.py | 42 ++++++++++++++++++++++-----------------
 2 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/nitric/api/storage.py b/nitric/api/storage.py
index 790f663..7e7f4b2 100644
--- a/nitric/api/storage.py
+++ b/nitric/api/storage.py
@@ -98,7 +98,7 @@ async def write(self, body: bytes):
     async def read(self) -> bytes:
         """Read this files contents from the bucket."""
         try:
-            response = await self._storage._storage_stub.read(StorageReadRequest(bucket_name=self._bucket, key=self.key))
+            response = await self._storage._storage_stub.read(storage_read_request=StorageReadRequest(bucket_name=self._bucket, key=self.key))
             return response.body
         except GRPCError as grpc_err:
             raise exception_from_grpc_error(grpc_err)
@@ -106,7 +106,7 @@ async def read(self) -> bytes:
     async def delete(self):
         """Delete this file from the bucket."""
         try:
-            await self._storage._storage_stub.delete(StorageDeleteRequest(bucket_name=self._bucket, key=self.key))
+            await self._storage._storage_stub.delete(storage_delete_request=StorageDeleteRequest(bucket_name=self._bucket, key=self.key))
         except GRPCError as grpc_err:
             raise exception_from_grpc_error(grpc_err)
 
diff --git a/tests/api/test_storage.py b/tests/api/test_storage.py
index a9e45bc..0124dc5 100644
--- a/tests/api/test_storage.py
+++ b/tests/api/test_storage.py
@@ -21,6 +21,8 @@
 
 import pytest
 from grpclib import GRPCError, Status
+from nitricapi.nitric.storage.v1 import StorageWriteRequest, StorageReadRequest, StorageDeleteRequest, \
+    StoragePreSignUrlRequest, StoragePreSignUrlRequestOperation
 
 from nitric.api import Storage
 from nitric.api.exception import UnknownException
@@ -44,10 +46,11 @@ async def test_write(self):
             await file.write(contents)
 
         # Check expected values were passed to Stub
-        mock_write.assert_called_once()
-        assert mock_write.call_args.kwargs["bucket_name"] == "test-bucket"
-        assert mock_write.call_args.kwargs["key"] == "test-file"
-        assert mock_write.call_args.kwargs["body"] == contents
+        mock_write.assert_called_once_with(storage_write_request=StorageWriteRequest(
+            bucket_name="test-bucket",
+            key="test-file",
+            body=contents
+        ))
 
     async def test_read(self):
         contents = b"some text as bytes"
@@ -65,23 +68,25 @@ async def test_read(self):
         assert response == contents
 
         # Check expected values were passed to Stub
-        mock_read.assert_called_once()
-        assert mock_read.call_args.kwargs["bucket_name"] == "test-bucket"
-        assert mock_read.call_args.kwargs["key"] == "test-file"
+        mock_read.assert_called_once_with(storage_read_request=StorageReadRequest(
+            bucket_name="test-bucket",
+            key="test-file",
+        ))
 
     async def test_delete(self):
-        mock_read = AsyncMock()
-        mock_read.return_value = Object()
+        mock_delete = AsyncMock()
+        mock_delete.return_value = Object()
 
-        with patch("nitricapi.nitric.storage.v1.StorageServiceStub.delete", mock_read):
+        with patch("nitricapi.nitric.storage.v1.StorageServiceStub.delete", mock_delete):
             bucket = Storage().bucket("test-bucket")
             file = bucket.file("test-file")
             await file.delete()
 
         # Check expected values were passed to Stub
-        mock_read.assert_called_once()
-        assert mock_read.call_args.kwargs["bucket_name"] == "test-bucket"
-        assert mock_read.call_args.kwargs["key"] == "test-file"
+        mock_delete.assert_called_once_with(storage_delete_request=StorageDeleteRequest(
+            bucket_name="test-bucket",
+            key="test-file",
+        ))
 
     async def test_sign_url(self):
         mock_pre_sign_url = AsyncMock()
@@ -93,11 +98,12 @@ async def test_sign_url(self):
             await file.sign_url()
 
         # Check expected values were passed to Stub
-        mock_pre_sign_url.assert_called_once()
-        assert mock_pre_sign_url.call_args.kwargs["bucket_name"] == "test-bucket"
-        assert mock_pre_sign_url.call_args.kwargs["key"] == "test-file"
-        assert mock_pre_sign_url.call_args.kwargs["operation"] == 0
-        assert mock_pre_sign_url.call_args.kwargs["expiry"] == 3600
+        mock_pre_sign_url.assert_called_once_with(storage_pre_sign_url_request=StoragePreSignUrlRequest(
+            bucket_name="test-bucket",
+            key="test-file",
+            operation=StoragePreSignUrlRequestOperation.READ,
+            expiry=3600
+        ))
 
     async def test_write_error(self):
         mock_write = AsyncMock()