From e624178b44b0112eff2f6773b62d01ea87c66bf8 Mon Sep 17 00:00:00 2001 From: Chris Shin Date: Wed, 9 Dec 2020 11:03:08 -0800 Subject: [PATCH] Fixes issue #754 by moving file read logic inside generator --- .../server/endpoint/fileuploads_endpoint.py | 16 +++-- test/assets/fileupload_append.xml | 3 + test/assets/fileupload_initialize.xml | 3 + test/test_fileuploads.py | 70 +++++++++++++++++++ 4 files changed, 86 insertions(+), 6 deletions(-) create mode 100644 test/assets/fileupload_append.xml create mode 100644 test/assets/fileupload_initialize.xml create mode 100644 test/test_fileuploads.py diff --git a/tableauserverclient/server/endpoint/fileuploads_endpoint.py b/tableauserverclient/server/endpoint/fileuploads_endpoint.py index 62224c894..c89a595d4 100644 --- a/tableauserverclient/server/endpoint/fileuploads_endpoint.py +++ b/tableauserverclient/server/endpoint/fileuploads_endpoint.py @@ -40,10 +40,18 @@ def append(self, xml_request, content_type): return FileuploadItem.from_response(server_response.content, self.parent_srv.namespace) def read_chunks(self, file): + file_opened = False + try: + file_content = open(file, 'rb') + file_opened = True + except TypeError: + file_content = file while True: - chunked_content = file.read(CHUNK_SIZE) + chunked_content = file_content.read(CHUNK_SIZE) if not chunked_content: + if file_opened: + file_content.close() break yield chunked_content @@ -52,11 +60,7 @@ def upload_chunks(cls, parent_srv, file): file_uploader = cls(parent_srv) upload_id = file_uploader.initiate() - try: - with open(file, 'rb') as f: - chunks = file_uploader.read_chunks(f) - except TypeError: - chunks = file_uploader.read_chunks(file) + chunks = file_uploader.read_chunks(file) for chunk in chunks: xml_request, content_type = RequestFactory.Fileupload.chunk_req(chunk) fileupload_item = file_uploader.append(xml_request, content_type) diff --git a/test/assets/fileupload_append.xml b/test/assets/fileupload_append.xml new file mode 100644 index 000000000..325ee66a9 --- /dev/null +++ b/test/assets/fileupload_append.xml @@ -0,0 +1,3 @@ + + + \ No newline at end of file diff --git a/test/assets/fileupload_initialize.xml b/test/assets/fileupload_initialize.xml new file mode 100644 index 000000000..073ad0edc --- /dev/null +++ b/test/assets/fileupload_initialize.xml @@ -0,0 +1,3 @@ + + + \ No newline at end of file diff --git a/test/test_fileuploads.py b/test/test_fileuploads.py new file mode 100644 index 000000000..9d115636f --- /dev/null +++ b/test/test_fileuploads.py @@ -0,0 +1,70 @@ +import os +import requests_mock +import unittest + +from ._utils import asset +from tableauserverclient.server import Server +from tableauserverclient.server.endpoint.fileuploads_endpoint import Fileuploads + +TEST_ASSET_DIR = os.path.join(os.path.dirname(__file__), 'assets') +FILEUPLOAD_INITIALIZE = os.path.join(TEST_ASSET_DIR, 'fileupload_initialize.xml') +FILEUPLOAD_APPEND = os.path.join(TEST_ASSET_DIR, 'fileupload_append.xml') + + +class FileuploadsTests(unittest.TestCase): + def setUp(self): + self.server = Server('http://test') + + # Fake sign in + self.server._site_id = 'dad65087-b08b-4603-af4e-2887b8aafc67' + self.server._auth_token = 'j80k54ll2lfMZ0tv97mlPvvSCRyD0DOM' + + self.baseurl = '{}/sites/{}/fileUploads'.format(self.server.baseurl, self.server.site_id) + + def test_read_chunks_file_path(self): + fileuploads = Fileuploads(self.server) + + file_path = asset('SampleWB.twbx') + chunks = fileuploads.read_chunks(file_path) + for chunk in chunks: + self.assertIsNotNone(chunk) + + def test_read_chunks_file_object(self): + fileuploads = Fileuploads(self.server) + + with open(asset('SampleWB.twbx'), 'rb') as f: + chunks = fileuploads.read_chunks(f) + for chunk in chunks: + self.assertIsNotNone(chunk) + + def test_upload_chunks_file_path(self): + fileuploads = Fileuploads(self.server) + file_path = asset('SampleWB.twbx') + upload_id = '7720:170fe6b1c1c7422dadff20f944d58a52-1:0' + + with open(FILEUPLOAD_INITIALIZE, 'rb') as f: + initialize_response_xml = f.read().decode('utf-8') + with open(FILEUPLOAD_APPEND, 'rb') as f: + append_response_xml = f.read().decode('utf-8') + with requests_mock.mock() as m: + m.post(self.baseurl, text=initialize_response_xml) + m.put(self.baseurl + '/' + upload_id, text=append_response_xml) + actual = fileuploads.upload_chunks(self.server, file_path) + + self.assertEqual(upload_id, actual) + + def test_upload_chunks_file_object(self): + fileuploads = Fileuploads(self.server) + upload_id = '7720:170fe6b1c1c7422dadff20f944d58a52-1:0' + + with open(asset('SampleWB.twbx'), 'rb') as file_content: + with open(FILEUPLOAD_INITIALIZE, 'rb') as f: + initialize_response_xml = f.read().decode('utf-8') + with open(FILEUPLOAD_APPEND, 'rb') as f: + append_response_xml = f.read().decode('utf-8') + with requests_mock.mock() as m: + m.post(self.baseurl, text=initialize_response_xml) + m.put(self.baseurl + '/' + upload_id, text=append_response_xml) + actual = fileuploads.upload_chunks(self.server, file_content) + + self.assertEqual(upload_id, actual)