Merge pull request #33 from box/add-preflight-option

Add preflight option for updating files

boxsdk/object/file.py

@@ -60,7 +60,7 @@ def download_to(self, writeable_stream):
         for chunk in box_response.network_response.response_as_stream.stream(decode_content=True):
             writeable_stream.write(chunk)
 
-    def update_contents_with_stream(self, file_stream, etag=None):
+    def update_contents_with_stream(self, file_stream, etag=None, preflight_check=False, preflight_expected_size=0):
         """
         Upload a new version of a file, taking the contents from the given file stream.
 
@@ -72,12 +72,26 @@ def update_contents_with_stream(self, file_stream, etag=None):
             If specified, instruct the Box API to update the item only if the current version's etag matches.
         :type etag:
             `unicode` or None
+        :param preflight_check:
+            If specified, preflight check will be performed before actually uploading the file.
+        :type preflight_check:
+            `bool`
+        :param preflight_expected_size:
+            The size of the file to be uploaded in bytes, which is used for preflight check. The default value is '0',
+            which means the file size is unknown.
+        :type preflight_expected_size:
+            `int`
         :returns:
             A new file object
         :rtype:
             :class:`File`
-        :raises: :class:`BoxAPIException` if the specified etag doesn't match the latest version of the file.
+        :raises:
+            :class:`BoxAPIException` if the specified etag doesn't match the latest version of the file or preflight
+            check fails.
         """
+        if preflight_check:
+            self.preflight_check(size=preflight_expected_size)
+
         url = self.get_url('content').replace(API.BASE_API_URL, API.UPLOAD_URL)
         files = {'file': ('unused', file_stream)}
         headers = {'If-Match': etag} if etag is not None else None
@@ -87,7 +101,7 @@ def update_contents_with_stream(self, file_stream, etag=None):
             response_object=self._session.post(url, expect_json_response=False, files=files, headers=headers).json(),
         )
 
-    def update_contents(self, file_path, etag=None):
+    def update_contents(self, file_path, etag=None, preflight_check=False, preflight_expected_size=0):
         """Upload a new version of a file. The contents are taken from the given file path.
 
         :param file_path:
@@ -98,14 +112,30 @@ def update_contents(self, file_path, etag=None):
             If specified, instruct the Box API to update the item only if the current version's etag matches.
         :type etag:
             `unicode` or None
+        :param preflight_check:
+            If specified, preflight check will be performed before actually uploading the file.
+        :type preflight_check:
+            `bool`
+        :param preflight_expected_size:
+            The size of the file to be uploaded in bytes, which is used for preflight check. The default value is '0',
+            which means the file size is unknown.
+        :type preflight_expected_size:
+            `int`
         :returns:
             A new file object
         :rtype:
             :class:`File`
-        :raises: :class:`BoxAPIException` if the specified etag doesn't match the latest version of the file.
+        :raises:
+            :class:`BoxAPIException` if the specified etag doesn't match the latest version of the file or preflight
+            check fails.
         """
         with open(file_path, 'rb') as file_stream:
-            return self.update_contents_with_stream(file_stream, etag)
+            return self.update_contents_with_stream(
+                file_stream,
+                etag,
+                preflight_check,
+                preflight_expected_size=preflight_expected_size,
+            )
 
     def lock(self, prevent_download=False):
         """

boxsdk/object/folder.py

@@ -143,7 +143,7 @@ def get_items(self, limit, offset=0, fields=None):
         response = box_response.json()
         return [Translator().translate(item['type'])(self._session, item['id'], item) for item in response['entries']]
 
-    def upload_stream(self, file_stream, file_name):
+    def upload_stream(self, file_stream, file_name, preflight_check=False, preflight_expected_size=0):
         """
         Upload a file to the folder.
         The contents are taken from the given file stream, and it will have the given name.
@@ -156,11 +156,23 @@ def upload_stream(self, file_stream, file_name):
             The name to give the file on Box.
         :type file_name:
             `unicode`
+        :param preflight_check:
+            If specified, preflight check will be performed before actually uploading the file.
+        :type preflight_check:
+            `bool`
+        :param preflight_expected_size:
+            The size of the file to be uploaded in bytes, which is used for preflight check. The default value is '0',
+            which means the file size is unknown.
+        :type preflight_expected_size:
+            `int`
         :returns:
             The newly uploaded file.
         :rtype:
             :class:`File`
         """
+        if preflight_check:
+            self.preflight_check(size=preflight_expected_size, name=file_name)
+
         url = '{0}/files/content'.format(API.UPLOAD_URL)
         data = {'attributes': json.dumps({
             'name': file_name,
@@ -178,7 +190,7 @@ def upload_stream(self, file_stream, file_name):
             response_object=file_response,
         )
 
-    def upload(self, file_path=None, file_name=None):
+    def upload(self, file_path=None, file_name=None, preflight_check=False, preflight_expected_size=0):
         """
         Upload a file to the folder.
         The contents are taken from the given file path, and it will have the given name.
@@ -192,6 +204,15 @@ def upload(self, file_path=None, file_name=None):
             The name to give the file on Box. If None, then use the leaf name of file_path
         :type file_name:
             `unicode`
+        :param preflight_check:
+            If specified, preflight check will be performed before actually uploading the file.
+        :type preflight_check:
+            `bool`
+        :param preflight_expected_size:
+            The size of the file to be uploaded in bytes, which is used for preflight check. The default value is '0',
+            which means the file size is unknown.
+        :type preflight_expected_size:
+            `int`
         :returns:
             The newly uploaded file.
         :rtype:
@@ -200,7 +221,12 @@ def upload(self, file_path=None, file_name=None):
         if file_name is None:
             file_name = os.path.basename(file_path)
         with open(file_path, 'rb') as file_stream:
-            return self.upload_stream(file_stream, file_name)
+            return self.upload_stream(
+                file_stream,
+                file_name,
+                preflight_check,
+                preflight_expected_size=preflight_expected_size,
+            )
 
     def create_subfolder(self, name):
         """

test/unit/object/conftest.py

@@ -144,6 +144,22 @@ def etag(request):
     return request.param
 
 
+@pytest.fixture(params=[True, False])
+def preflight_check(request):
+    return request.param
+
+
+@pytest.fixture(params=[True, False])
+def preflight_fails(preflight_check, request):
+    # pylint:disable=redefined-outer-name
+    return preflight_check and request.param
+
+
+@pytest.fixture(params=[0, 100])
+def file_size(request):
+    return request.param
+
+
 @pytest.fixture()
 def mock_group(mock_box_session, mock_group_id):
     # pylint:disable=redefined-outer-name

test/unit/object/test_file.py

@@ -6,7 +6,7 @@
 import pytest
 from six import BytesIO
 from boxsdk.config import API
-
+from boxsdk.exception import BoxAPIException
 from boxsdk.object.file import File
 
 
@@ -76,6 +76,66 @@ def test_update_content(
     assert new_file.object_id == mock_upload_response.json()['entries'][0]['id']
 
 
+def test_update_contents_with_stream_does_preflight_check_if_specified(
+        test_file,
+        preflight_check,
+        file_size,
+        preflight_fails,
+        mock_box_session,
+):
+    with patch.object(File, 'preflight_check', return_value=None):
+        kwargs = {'file_stream': BytesIO(b'some bytes')}
+        if preflight_check:
+            kwargs['preflight_check'] = preflight_check
+            kwargs['preflight_expected_size'] = file_size
+        if preflight_fails:
+            test_file.preflight_check.side_effect = BoxAPIException(400)
+            with pytest.raises(BoxAPIException):
+                test_file.update_contents_with_stream(**kwargs)
+        else:
+            test_file.update_contents_with_stream(**kwargs)
+
+        if preflight_check:
+            assert test_file.preflight_check.called_once_with(size=file_size)
+            if preflight_fails:
+                assert not mock_box_session.post.called
+            else:
+                assert mock_box_session.post.called
+        else:
+            assert not test_file.preflight_check.called
+
+
+@patch('boxsdk.object.file.open', mock_open(read_data=b'some bytes'), create=True)
+def test_update_contents_does_preflight_check_if_specified(
+        test_file,
+        mock_file_path,
+        preflight_check,
+        file_size,
+        preflight_fails,
+        mock_box_session,
+):
+    with patch.object(File, 'preflight_check', return_value=None):
+        kwargs = {'file_path': mock_file_path}
+        if preflight_check:
+            kwargs['preflight_check'] = preflight_check
+            kwargs['preflight_expected_size'] = file_size
+        if preflight_fails:
+            test_file.preflight_check.side_effect = BoxAPIException(400)
+            with pytest.raises(BoxAPIException):
+                test_file.update_contents(**kwargs)
+        else:
+            test_file.update_contents(**kwargs)
+
+        if preflight_check:
+            assert test_file.preflight_check.called_once_with(size=file_size)
+            if preflight_fails:
+                assert not mock_box_session.post.called
+            else:
+                assert mock_box_session.post.called
+        else:
+            assert not test_file.preflight_check.called
+
+
 @pytest.mark.parametrize('prevent_download', (True, False))
 def test_lock(test_file, mock_box_session, mock_file_response, prevent_download):
     expected_url = test_file.get_url()

test/unit/object/test_folder.py

@@ -3,11 +3,12 @@
 from __future__ import unicode_literals
 import json
 from os.path import basename
-from mock import mock_open, patch, Mock
+from mock import mock_open, patch, Mock, MagicMock
 import pytest
 from six import BytesIO
 from six.moves import zip  # pylint:disable=redefined-builtin,import-error
 from boxsdk.config import API
+from boxsdk.exception import BoxAPIException
 from boxsdk.network.default_network import DefaultNetworkResponse
 from boxsdk.object.file import File
 from boxsdk.object.collaboration import Collaboration, CollaborationRole
@@ -148,6 +149,69 @@ def test_upload(
     assert new_file.object_id == mock_upload_response.json()['entries'][0]['id']
 
 
+def test_upload_stream_does_preflight_check_if_specified(
+        mock_box_session,
+        test_folder,
+        preflight_check,
+        preflight_fails,
+        file_size,
+):
+    with patch.object(Folder, 'preflight_check', return_value=None):
+        kwargs = {'file_stream': BytesIO(b'some bytes'), 'file_name': 'foo.txt'}
+        mock_box_session.post = MagicMock()
+        if preflight_check:
+            kwargs['preflight_check'] = preflight_check
+            kwargs['preflight_expected_size'] = file_size
+        if preflight_fails:
+            test_folder.preflight_check.side_effect = BoxAPIException(400)
+            with pytest.raises(BoxAPIException):
+                test_folder.upload_stream(**kwargs)
+        else:
+            test_folder.upload_stream(**kwargs)
+
+        if preflight_check:
+            assert test_folder.preflight_check.called_once_with(size=file_size, name='foo.txt')
+            _assert_post_called_correctly(mock_box_session, preflight_fails)
+        else:
+            assert not test_folder.preflight_check.called
+
+
+def _assert_post_called_correctly(mock_box_session, preflight_fails):
+    if preflight_fails:
+        assert not mock_box_session.post.called
+    else:
+        assert mock_box_session.post.called
+
+
+@patch('boxsdk.object.folder.open', mock_open(read_data=b'some bytes'), create=True)
+def test_upload_does_preflight_check_if_specified(
+        mock_box_session,
+        test_folder,
+        mock_file_path,
+        preflight_check,
+        preflight_fails,
+        file_size,
+):
+    with patch.object(Folder, 'preflight_check', return_value=None):
+        kwargs = {'file_path': mock_file_path, 'file_name': 'foo.txt'}
+        mock_box_session.post = MagicMock()
+        if preflight_check:
+            kwargs['preflight_check'] = preflight_check
+            kwargs['preflight_expected_size'] = file_size
+        if preflight_fails:
+            test_folder.preflight_check.side_effect = BoxAPIException(400)
+            with pytest.raises(BoxAPIException):
+                test_folder.upload(**kwargs)
+        else:
+            test_folder.upload(**kwargs)
+
+        if preflight_check:
+            assert test_folder.preflight_check.called_once_with(size=file_size, name='foo.txt')
+            _assert_post_called_correctly(mock_box_session, preflight_fails)
+        else:
+            assert not test_folder.preflight_check.called
+
+
 def test_create_subfolder(test_folder, mock_box_session, mock_object_id, mock_folder_response):
     expected_url = test_folder.get_type_url()
     mock_box_session.post.return_value = mock_folder_response