diff --git a/grizzly/users/blobstorage.py b/grizzly/users/blobstorage.py index 01551b95..14fca352 100644 --- a/grizzly/users/blobstorage.py +++ b/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 diff --git a/tests/unit/test_grizzly/users/test_blobstorage.py b/tests/unit/test_grizzly/users/test_blobstorage.py index 4f1fb56b..16b3a50d 100644 --- a/tests/unit/test_grizzly/users/test_blobstorage.py +++ b/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()