Spaces:
Running on Zero
Running on Zero
File size: 2,326 Bytes
28a3312 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | import os
import tempfile
import unittest
from unittest.mock import MagicMock, patch
from PIL import Image
import drive_storage
class DriveStorageTests(unittest.TestCase):
def test_save_is_disabled_without_credentials(self):
with patch.dict(os.environ, {}, clear=True):
self.assertFalse(drive_storage.save_images_async("missing.jpg", None, 0))
def test_upload_retries_transient_failure(self):
failed = MagicMock(ok=False, status_code=503)
succeeded = MagicMock(ok=True)
succeeded.json.return_value = {"id": "file-id"}
session = MagicMock()
session.post.side_effect = [failed, succeeded]
with patch("drive_storage.time.sleep"):
file_id = drive_storage._upload_image(
session,
"token",
"image.jpg",
b"jpeg",
{"source": "crowd-counter"},
)
self.assertEqual(file_id, "file-id")
self.assertEqual(session.post.call_count, 2)
def test_save_prepares_original_and_overlay(self):
submitted = []
original_executor = drive_storage._executor
drive_storage._executor = MagicMock()
drive_storage._executor.submit.side_effect = lambda *args: submitted.append(args)
try:
with tempfile.NamedTemporaryFile(suffix=".png") as image_file:
Image.new("RGB", (8, 8), "white").save(image_file.name)
with patch.dict(
os.environ,
{"GOOGLE_SERVICE_ACCOUNT_JSON": '{"client_email":"bot@example.com"}'},
clear=True,
):
saved = drive_storage.save_images_async(
image_file.name,
Image.new("RGB", (8, 8), "black"),
42,
)
finally:
drive_storage._executor = original_executor
if submitted:
drive_storage._pending_saves.release()
self.assertTrue(saved)
self.assertEqual(len(submitted), 1)
self.assertEqual(submitted[0][-1], 42)
self.assertTrue(submitted[0][-3].startswith(b"\xff\xd8"))
self.assertTrue(submitted[0][-2].startswith(b"\xff\xd8"))
if __name__ == "__main__":
unittest.main()
|