diff --git a/tests/core/test_remote_logging.py b/tests/core/test_remote_logging.py index 88ca8bd42..cdfb55fad 100644 --- a/tests/core/test_remote_logging.py +++ b/tests/core/test_remote_logging.py @@ -3,6 +3,77 @@ from waterbutler.core import remote_logging +class TestLogToCallback: + + @pytest.mark.asyncio + @pytest.mark.parametrize('completed, expected', [(False, False), (True, True)]) + async def test_download_action_sets_completed_flag(self, monkeypatch, completed, expected): + captured = {} + + async def fake_send_signed_request(method, url, payload): + captured['payload'] = payload + return 200, b'success' + + monkeypatch.setattr(remote_logging.utils, 'send_signed_request', fake_send_signed_request) + + class DummySource: + auth = {'callback_url': 'https://example.com/callback'} + + def serialize(self): + return {'provider': 'osf'} + + source = DummySource() + request = { + 'request': { + 'method': 'GET', + 'url': 'https://example.com/file', + 'headers': {}, + }, + 'referrer': {'url': None}, + 'tech': {'ua': 'test-agent', 'ip': '127.0.0.1'}, + } + + await remote_logging.log_to_callback( + 'download_file', + source=source, + request=request, + completed=completed, + ) + + assert captured['payload']['action_meta']['completed'] is expected + + @pytest.mark.asyncio + async def test_non_download_action_omits_completed_flag(self, monkeypatch): + captured = {} + + async def fake_send_signed_request(method, url, payload): + captured['payload'] = payload + return 200, b'success' + + monkeypatch.setattr(remote_logging.utils, 'send_signed_request', fake_send_signed_request) + + class DummySource: + auth = {'callback_url': 'https://example.com/callback'} + + def serialize(self): + return {'provider': 'osf'} + + source = DummySource() + request = { + 'request': { + 'method': 'GET', + 'url': 'https://example.com/file', + 'headers': {}, + }, + 'referrer': {'url': None}, + 'tech': {'ua': 'test-agent', 'ip': '127.0.0.1'}, + } + + await remote_logging.log_to_callback('create', source=source, request=request) + + assert 'completed' not in captured['payload']['action_meta'] + + class TestScrubPayloadForKeen: def test_flat_dict(self): diff --git a/tests/server/api/v1/test_metadata_mixin.py b/tests/server/api/v1/test_metadata_mixin.py index 48308e806..09aafabb1 100644 --- a/tests/server/api/v1/test_metadata_mixin.py +++ b/tests/server/api/v1/test_metadata_mixin.py @@ -120,6 +120,20 @@ async def test_download_file_headers_no_stream_name(self, http_request, mock_str handler.write_stream.assert_awaited_once() + @pytest.mark.asyncio + @pytest.mark.parametrize('write_stream_result, expected_completed', [(True, True), (False, False)]) + async def test_download_file_records_stream_completion(self, http_request, mock_stream, + write_stream_result, expected_completed): + + handler = mock_handler(http_request) + handler.provider.download = MockCoroutine(return_value=mock_stream) + handler.path = WaterButlerPath('/test_file') + handler.write_stream = MockCoroutine(return_value=write_stream_result) + + await handler.download_file() + + assert handler._download_completed is expected_completed + @pytest.mark.asyncio @pytest.mark.parametrize("given_arg,expected_name,filtered_name", [ (['résumé.doc'], 'r%C3%A9sum%C3%A9.doc', 'resume.doc'), diff --git a/tests/server/api/v1/test_provider.py b/tests/server/api/v1/test_provider.py index a39eb9911..e9445635d 100644 --- a/tests/server/api/v1/test_provider.py +++ b/tests/server/api/v1/test_provider.py @@ -165,15 +165,17 @@ async def test_data_received_stream(self, http_request): class TestProviderHandlerFinish: @pytest.mark.asyncio - async def test_on_finish_download_file(self, http_request): + @pytest.mark.parametrize('download_completed, expected_completed', [(True, True), (False, False)]) + async def test_on_finish_download_file(self, http_request, download_completed, expected_completed): handler = mock_handler(http_request) handler.request.method = 'GET' handler.path = WaterButlerPath('/file') + handler._download_completed = download_completed handler._send_hook = mock.Mock() assert handler.on_finish() is None - handler._send_hook.assert_called_once_with('download_file') + handler._send_hook.assert_called_once_with('download_file', completed=expected_completed) @pytest.mark.asyncio async def test_on_finish_download_zip(self, http_request): diff --git a/waterbutler/core/remote_logging.py b/waterbutler/core/remote_logging.py index bc2d256cf..9debe58c4 100644 --- a/waterbutler/core/remote_logging.py +++ b/waterbutler/core/remote_logging.py @@ -18,7 +18,7 @@ @utils.async_retry(retries=5, backoff=5) async def log_to_callback(action, source=None, destination=None, start_time=None, errors=None, - request=None): + request=None, bytes_downloaded=0, completed=False): """PUT a logging payload back to the callback given by the auth provider.""" errors = errors or [] request = request or {} @@ -59,7 +59,10 @@ async def log_to_callback(action, source=None, destination=None, start_time=None is_mfr_render = (ref_url_domain == settings.MFR_DOMAIN or settings.MFR_IDENTIFYING_HEADER in request["request"]["headers"]) log_payload['action_meta']['is_mfr_render'] = is_mfr_render + log_payload['action_meta']['completed'] = completed + log_payload['action_meta']['bytes_downloaded'] = bytes_downloaded + log_payload['action_meta']['ip'] = request['tech']['ip'] resp_status, resp_data = await utils.send_signed_request('PUT', auth['callback_url'], log_payload) if resp_status // 100 != 2: @@ -217,12 +220,14 @@ async def _send_to_keen(payload, collection, project_id, write_key, action, doma def log_file_action(action, source, api_version, destination=None, request=None, - start_time=None, errors=None, bytes_downloaded=None, bytes_uploaded=None): + start_time=None, errors=None, bytes_downloaded=None, bytes_uploaded=None, + completed=False): """Kick off logging actions in the background. Returns array of asyncio.Tasks.""" request = request or {} return [ log_to_callback(action, source=source, destination=destination, - start_time=start_time, errors=errors, request=request,), + start_time=start_time, errors=errors, request=request, + bytes_downloaded=bytes_downloaded, completed=completed,), asyncio.ensure_future( log_to_keen(action, source=source, destination=destination, errors=errors, request=request, api_version=api_version, diff --git a/waterbutler/server/api/v0/core.py b/waterbutler/server/api/v0/core.py index 543d966a2..af702e186 100644 --- a/waterbutler/server/api/v0/core.py +++ b/waterbutler/server/api/v0/core.py @@ -94,12 +94,13 @@ async def prepare(self): self.path = await self.provider.validate_path(**self.arguments) self.arguments['path'] = self.path # TODO Not this - def _send_hook(self, action, metadata=None, path=None): + def _send_hook(self, action, metadata=None, path=None, completed=False): source = LogPayload(self.arguments['nid'], self.provider, metadata=metadata, path=path) remote_logging.log_file_action(action, source=source, api_version='v0', request=remote_logging._serialize_request(self.request), bytes_downloaded=self.bytes_downloaded, - bytes_uploaded=self.bytes_uploaded) + bytes_uploaded=self.bytes_uploaded, + completed=completed) class BaseCrossProviderHandler(BaseHandler): @@ -140,7 +141,7 @@ def json(self): return self._json - def _send_hook(self, action, metadata): + def _send_hook(self, action, metadata, completed=False): source = LogPayload(self.json['source']['nid'], self.source_provider, path=self.json['source']['path']) destination = LogPayload(self.json['destination']['nid'], self.destination_provider, @@ -148,4 +149,5 @@ def _send_hook(self, action, metadata): remote_logging.log_file_action(action, source=source, destination=destination, api_version='v0', request=remote_logging._serialize_request(self.request), bytes_downloaded=self.bytes_downloaded, - bytes_uploaded=self.bytes_uploaded) + bytes_uploaded=self.bytes_uploaded, + completed=completed) diff --git a/waterbutler/server/api/v0/crud.py b/waterbutler/server/api/v0/crud.py index 07034155f..0957d0042 100644 --- a/waterbutler/server/api/v0/crud.py +++ b/waterbutler/server/api/v0/crud.py @@ -61,7 +61,7 @@ async def get(self): if isinstance(result, str): self.redirect(result) - self._send_hook('download_file', path=self.path) + self._send_hook('download_file', path=self.path, completed=True) return if getattr(result, 'partial', None): @@ -86,8 +86,8 @@ async def get(self): if ext in mime_types: self.set_header('Content-Type', mime_types[ext]) - await self.write_stream(result) - self._send_hook('download_file', path=self.path) + completed = await self.write_stream(result) + self._send_hook('download_file', path=self.path, completed=completed) async def post(self): """Create a folder""" diff --git a/waterbutler/server/api/v0/zip.py b/waterbutler/server/api/v0/zip.py index 2ed0695c6..34c64a744 100644 --- a/waterbutler/server/api/v0/zip.py +++ b/waterbutler/server/api/v0/zip.py @@ -20,5 +20,5 @@ async def get(self): result = await self.provider.zip(**self.arguments) - await self.write_stream(result) - self._send_hook('download_zip', path=self.path) + completed = await self.write_stream(result) + self._send_hook('download_zip', path=self.path, completed=completed) diff --git a/waterbutler/server/api/v1/provider/__init__.py b/waterbutler/server/api/v1/provider/__init__.py index 3cf8ed318..a0defc6fe 100644 --- a/waterbutler/server/api/v1/provider/__init__.py +++ b/waterbutler/server/api/v1/provider/__init__.py @@ -246,6 +246,7 @@ def on_finish(self): 'zip' not in self.request.query_arguments))): return + completed = False # Done here just because method is defined action = { 'GET': lambda: 'download_file' if self.path.is_file else 'download_zip', @@ -254,9 +255,14 @@ def on_finish(self): 'DELETE': lambda: 'delete' }[method]() + if action in {'download_file', 'download_zip'}: + completed = getattr(self, '_download_completed', status in {200, 302}) + self._send_hook(action, completed=completed) + return + self._send_hook(action) - def _send_hook(self, action): + def _send_hook(self, action, completed=False): source = None destination = None @@ -281,4 +287,5 @@ def _send_hook(self, action): remote_logging.log_file_action(action, source=source, destination=destination, api_version='v1', request=remote_logging._serialize_request(self.request), bytes_downloaded=self.bytes_downloaded, - bytes_uploaded=self.bytes_uploaded,) + bytes_uploaded=self.bytes_uploaded, + completed=completed) diff --git a/waterbutler/server/api/v1/provider/metadata.py b/waterbutler/server/api/v1/provider/metadata.py index add78e150..988a91f66 100644 --- a/waterbutler/server/api/v1/provider/metadata.py +++ b/waterbutler/server/api/v1/provider/metadata.py @@ -73,6 +73,7 @@ async def download_file(self): ) if isinstance(stream, str): + self._download_completed = True return self.redirect(stream) if getattr(stream, 'partial', None): @@ -103,7 +104,7 @@ async def download_file(self): if ext in mime_types: self.set_header('Content-Type', mime_types[ext]) - await self.write_stream(stream) + self._download_completed = await self.write_stream(stream) if getattr(stream, 'partial', False) and isinstance(stream, ResponseStreamReader): await stream.response.release() @@ -133,4 +134,4 @@ async def download_folder_as_zip(self): result = await self.provider.zip(self.path) - await self.write_stream(result) + self._download_completed = await self.write_stream(result) diff --git a/waterbutler/server/utils.py b/waterbutler/server/utils.py index 83d2a01a2..e9afd243e 100644 --- a/waterbutler/server/utils.py +++ b/waterbutler/server/utils.py @@ -123,7 +123,6 @@ def set_status(self, code, reason=None): async def write_stream(self, stream): try: - while True: chunk = await stream.read(settings.CHUNK_SIZE) if not chunk: @@ -138,4 +137,6 @@ async def write_stream(self, stream): except tornado.iostream.StreamClosedError: # Client has disconnected early. # No need for any exception to be raised - return + return False + + return True