From 55000d5c483f9aded276652e50ad49bfafa50e3a Mon Sep 17 00:00:00 2001 From: Oliver Walters Date: Tue, 17 May 2022 22:28:46 +1000 Subject: [PATCH] Add ability to download file as StringIO or BytesIO --- InvenTree/InvenTree/api_tester.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/InvenTree/InvenTree/api_tester.py b/InvenTree/InvenTree/api_tester.py index 0420e93369..122ce1950f 100644 --- a/InvenTree/InvenTree/api_tester.py +++ b/InvenTree/InvenTree/api_tester.py @@ -170,7 +170,7 @@ class InvenTreeAPITestCase(APITestCase): return response - def download_file(self, url, data, expected_code=None, expected_fn=None): + def download_file(self, url, data, expected_code=None, expected_fn=None, decode=False): """ Download a file from the server, and return an in-memory file """ @@ -191,11 +191,20 @@ class InvenTreeAPITestCase(APITestCase): fn = result.groups()[0] - with io.BytesIO() as fo: - fo.name = fn - fo.write(response.getvalue()) - if expected_fn is not None: self.assertEqual(expected_fn, fn) - - return fo \ No newline at end of file + + if decode: + # Decode data and return as StringIO file object + fo = io.StringIO() + fo.name = fo + fo.write(response.getvalue().decode('UTF-8')) + else: + # Return a a BytesIO file object + fo = io.BytesIO() + fo.name = fn + fo.write(response.getvalue()) + + fo.seek(0) + + return fo