BlobStorageUser: support for RECEIVE/GET (#277)

* Implement downloading of blobs

* get metadata for blob and enable response handlers

* always overwrite if destination file exists when writing

* test coverage

grizzly/users/blobstorage.py

@@ -6,6 +6,8 @@
 
 * send
 * put
+* receive
+* put
 
 ## Format
 
@@ -34,15 +36,15 @@
 
 from azure.storage.blob import BlobServiceClient
 
-from grizzly.types import RequestMethod, GrizzlyResponse
+from grizzly.types import RequestMethod, GrizzlyResponse, RequestDirection
 from grizzly.types.locust import Environment
 from grizzly.tasks import RequestTask
 from grizzly.utils import merge_dicts
 
-from .base import GrizzlyUser
+from .base import GrizzlyUser, ResponseHandler
 
 
-class BlobStorageUser(GrizzlyUser):
+class BlobStorageUser(ResponseHandler, GrizzlyUser):
     blob_client: BlobServiceClient
     _context: Dict[str, Any] = {}
 
@@ -55,16 +57,13 @@ def __init__(self, environment: Environment, *args: Tuple[Any], **kwargs: Dict[s
 
         # Replace semicolon separators between parameters to ? and & and massage it to make it "urlparse-compliant"
         # for validation
-        conn_str = conn_str.replace(';EndpointSuffix=', '://', 1).replace(';', '/?', 1).replace(';', '&')
+        conn_str = conn_str.replace(';', '://?', 1).replace(';', '&')
 
         parsed = urlparse(conn_str)
 
         if parsed.scheme != 'https':
             raise ValueError(f'"{parsed.scheme}" is not supported for {self.__class__.__name__}')
 
-        if parsed.query == '':
-            raise ValueError(f'{self.__class__.__name__} needs AccountName and AccountKey in the query string')
-
         params = parse_qs(parsed.query)
         if 'AccountName' not in params:
             raise ValueError(f'{self.__class__.__name__} needs AccountName in the query string')
@@ -85,17 +84,23 @@ def on_stop(self) -> None:
     def request_impl(self, request: RequestTask) -> GrizzlyResponse:
         blob = os.path.basename(request.endpoint)
         container = request.endpoint
-        print(f'{request.endpoint=}, {container=}, {blob=}')
 
         if container.endswith(blob):
             container = os.path.dirname(container)
         else:
             blob = self.normalize(request.name)
 
+        if request.method not in [RequestMethod.SEND, RequestMethod.PUT, RequestMethod.RECEIVE, RequestMethod.GET]:
+            raise NotImplementedError(f'{self.__class__.__name__} has not implemented {request.method.name}')
+
         with self.blob_client.get_blob_client(container=container, blob=blob) as blob_client:
-            if request.method in [RequestMethod.SEND, RequestMethod.PUT]:
-                blob_client.upload_blob(request.source)
-            else:  # pragma: no cover
-                raise NotImplementedError(f'{self.__class__.__name__} has not implemented {request.method.name}')
+            if request.method.direction == RequestDirection.TO:
+                blob_client.upload_blob(request.source, overwrite=True)
+            else:
+                downloader = blob_client.download_blob()
+                request.source = downloader.readall().decode('utf-8')
+
+        properties = blob_client.get_blob_properties()
+        headers = {key: value for key, value in properties.items()}
 
-        return {}, request.source
+        return headers, request.source

tests/unit/test_grizzly/users/test_blobstorage.py

@@ -1,6 +1,7 @@
 import json
 
 from typing import cast
+from unittest.mock import ANY
 
 import pytest
 
@@ -81,17 +82,17 @@ def test_create(self, blob_storage_parent: GrizzlyScenario) -> None:
         BlobStorageUser.host = 'DefaultEndpointsProtocol=https;EndpointSuffix=core.windows.net'
         with pytest.raises(ValueError) as e:
             BlobStorageUser(blob_storage_parent.user.environment)
-        assert 'needs AccountName and AccountKey in the query string' in str(e)
+        assert 'BlobStorageUser needs AccountName in the query string' in str(e)
 
         BlobStorageUser.host = 'DefaultEndpointsProtocol=https;EndpointSuffix=core.windows.net;AccountKey=xxxyyyyzzz=='
         with pytest.raises(ValueError) as e:
             BlobStorageUser(blob_storage_parent.user.environment)
-        assert 'needs AccountName in the query string' in str(e)
+        assert 'BlobStorageUser needs AccountName in the query string' in str(e)
 
         BlobStorageUser.host = 'DefaultEndpointsProtocol=https;EndpointSuffix=core.windows.net;AccountName=my-storage'
         with pytest.raises(ValueError) as e:
             BlobStorageUser(blob_storage_parent.user.environment)
-        assert 'needs AccountKey in the query string' in str(e)
+        assert 'BlobStorageUser needs AccountKey in the query string' in str(e)
 
     @pytest.mark.usefixtures('blob_storage_parent')
     def test_send(self, blob_storage_parent: GrizzlyScenario, mocker: MockerFixture, grizzly_fixture: GrizzlyFixture) -> None:
@@ -111,6 +112,7 @@ def test_send(self, blob_storage_parent: GrizzlyScenario, mocker: MockerFixture,
         request.endpoint = 'some_container_name/file.txt'
 
         upload_blob = mocker.patch('azure.storage.blob._blob_service_client.BlobClient.upload_blob', autospec=True)
+        blob_properties = mocker.patch('azure.storage.blob._blob_service_client.BlobClient.get_blob_properties', return_value={'size': 1337, 'etag': '0xdeadbeef'})
 
         expected_payload = {
             'result': {
@@ -123,44 +125,135 @@ def test_send(self, blob_storage_parent: GrizzlyScenario, mocker: MockerFixture,
             }
         }
 
+        assert isinstance(blob_storage_parent.user, BlobStorageUser)
+
         metadata, payload = blob_storage_parent.user.request(request)
 
         assert payload is not None
-        assert upload_blob.call_count == 1
+        upload_blob.assert_called_once_with(ANY, json.dumps(expected_payload, indent=4), overwrite=True)
         args, _ = upload_blob.call_args_list[-1]
-        assert len(args) == 2
-        assert args[1] == json.dumps(expected_payload, indent=4)
 
         blob_client = cast(BlobClient, args[0])
 
-        assert metadata == {}
+        assert metadata == {'size': 1337, 'etag': '0xdeadbeef'}
 
         json_payload = json.loads(payload)
         assert json_payload['result']['id'] == 'ID-31337'
         assert blob_client.container_name == 'some_container_name'
         assert blob_client.blob_name == 'file.txt'
 
+        blob_properties.assert_called_once_with()
+        blob_properties.reset_mock()
+
+        upload_blob = mocker.patch('azure.storage.blob._blob_service_client.BlobClient.upload_blob', side_effect=[RuntimeError('failed to upload blob')])
+
+        request.method = RequestMethod.SEND
+        blob_storage_parent.user._scenario.failure_exception = RestartScenario
+
+        with pytest.raises(RestartScenario):
+            blob_storage_parent.user.request(request)
+
+        blob_properties.assert_not_called()
+
+    @pytest.mark.usefixtures('blob_storage_parent')
+    def test_receive(self, blob_storage_parent: GrizzlyScenario, mocker: MockerFixture, grizzly_fixture: GrizzlyFixture) -> None:
+        blob_storage_parent.user.on_start()
+
+        grizzly = grizzly_fixture.grizzly
+
+        remote_variables = {
+            'variables': transform(grizzly, {
+                'AtomicIntegerIncrementer.messageID': 31337,
+                'AtomicDate.now': '',
+                'messageID': 137,
+            }),
+        }
+        blob_storage_parent.user.add_context(remote_variables)
         request = cast(RequestTask, blob_storage_parent.user._scenario.tasks()[-1])
+        request.method = RequestMethod.RECEIVE
+        request.endpoint = 'some_container_name/file.txt'
+
+        download_blob = mocker.patch('azure.storage.blob._blob_service_client.BlobClient.download_blob', autospec=True)
+        blob_properties = mocker.patch('azure.storage.blob._blob_service_client.BlobClient.get_blob_properties', return_value={'size': 7331, 'etag': '0xb000b000'})
+
+        expected_payload = {
+            'result': {
+                'id': 'ID-31337',
+                'date': '',
+                'variable': '137',
+                'item': {
+                    'description': 'this is just a description',
+                }
+            }
+        }
+
+        download_blob.return_value.readall.return_value = json.dumps(expected_payload, indent=4).encode('utf-8')
+
+        metadata, payload = blob_storage_parent.user.request(request)
 
-        blob_storage_parent.user.request(request)
+        assert payload == json.dumps(expected_payload, indent=4)
+        args, _ = download_blob.call_args_list[-1]
 
+        blob_client = cast(BlobClient, args[0])
+        download_blob.assert_called_once_with(blob_client)
+        download_blob.return_value.readall.assert_called_once_with()
+
+        assert metadata == {'size': 7331, 'etag': '0xb000b000'}
+
+        json_payload = json.loads(payload)
+        assert json_payload['result']['id'] == 'ID-31337'
+        assert blob_client.container_name == 'some_container_name'
+        assert blob_client.blob_name == 'file.txt'
+
+        blob_properties.assert_called_once_with()
+        blob_properties.reset_mock()
+
+        download_blob = mocker.patch('azure.storage.blob._blob_service_client.BlobClient.download_blob', side_effect=[RuntimeError('failed to download blob')])
+
+        blob_storage_parent.user._scenario.failure_exception = StopUser
+
+        with pytest.raises(StopUser):
+            blob_storage_parent.user.request(request)
+
+        blob_properties.assert_not_called()
+
+    @pytest.mark.usefixtures('blob_storage_parent')
+    def test_not_implemented(self, blob_storage_parent: GrizzlyScenario, mocker: MockerFixture, grizzly_fixture: GrizzlyFixture) -> None:
+        blob_storage_parent.user.on_start()
+
+        grizzly = grizzly_fixture.grizzly
+
+        remote_variables = {
+            'variables': transform(grizzly, {
+                'AtomicIntegerIncrementer.messageID': 31337,
+                'AtomicDate.now': '',
+                'messageID': 137,
+            }),
+        }
+        blob_storage_parent.user.add_context(remote_variables)
+
+        request = cast(RequestTask, blob_storage_parent.user._scenario.tasks()[-1])
+        request.endpoint = 'some_container_name/file.txt'
+
+        download_blob = mocker.patch('azure.storage.blob._blob_service_client.BlobClient.download_blob', autospec=True)
+        upload_blob = mocker.patch('azure.storage.blob._blob_service_client.BlobClient.upload_blob', autospec=True)
         request_event = mocker.spy(blob_storage_parent.user.environment.events.request, 'fire')
 
-        request.method = RequestMethod.RECEIVE
+        request.method = RequestMethod.POST
         with pytest.raises(StopUser):
             blob_storage_parent.user.request(request)
 
         assert request_event.call_count == 1
         _, kwargs = request_event.call_args_list[-1]
 
-        assert kwargs.get('request_type', None) == 'RECV'
+        assert kwargs.get('request_type', None) == 'POST'
         assert kwargs.get('name', None) == f'{blob_storage_parent.user._scenario.identifier} {request.name}'
         assert kwargs.get('response_time', -1) > -1
         assert kwargs.get('response_length', None) == 0
         assert kwargs.get('context', None) is blob_storage_parent.user._context
         exception = kwargs.get('exception', None)
         assert isinstance(exception, NotImplementedError)
-        assert 'has not implemented RECEIVE' in str(exception)
+        assert 'has not implemented POST' in str(exception)
 
         blob_storage_parent.user._scenario.failure_exception = None
         with pytest.raises(StopUser):
@@ -174,9 +267,5 @@ def test_send(self, blob_storage_parent: GrizzlyScenario, mocker: MockerFixture,
         with pytest.raises(StopUser):
             blob_storage_parent.user.request(request)
 
-        upload_blob = mocker.patch('azure.storage.blob._blob_service_client.BlobClient.upload_blob', side_effect=[RuntimeError('failed to upload blob')])
-
-        request.method = RequestMethod.SEND
-
-        with pytest.raises(RestartScenario):
-            blob_storage_parent.user.request(request)
+        download_blob.assert_not_called()
+        upload_blob.assert_not_called()