import io import os from pathlib import Path import tarfile import tempfile import unittest from unittest.mock import patch import httpx import trimesh from labeler.cloud import GSOCloudSource, CloudDatasetError, unpack_archive from labeler.dataset import cached_measure def archive_for(asset_id, nested=False): output = io.BytesIO() mesh = trimesh.creation.box([3, 2, 1]).export(file_type='obj').encode() with tarfile.open(fileobj=output, mode='w:gz') as archive: member = tarfile.TarInfo(f'{asset_id}/visual_geometry.obj' if nested else 'visual_geometry.obj') member.size = len(mesh) archive.addfile(member, io.BytesIO(mesh)) return output.getvalue() class CloudTests(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.addCleanup(self.temp.cleanup) self.root = Path(self.temp.name) self.requests = [] self.manifest = {'version': 1, 'data_dir': 'gs://test-bucket/GSO', 'assets': { name: {'asset_type': 'FileBasedObject', 'id': name, 'license': 'CC BY-SA 4.0', 'kwargs': {'render_filename': '{asset_dir}/visual_geometry.obj'}} for name in ['one', 'two']}} self.fail_archives = False def respond(request): self.requests.append(str(request.url)) if request.url.path.endswith('.json'): return httpx.Response(200, json=self.manifest) if self.fail_archives: return httpx.Response(503) name = request.url.path.rsplit('/', 1)[-1].split('.')[0] return httpx.Response(200, content=archive_for(name, nested=name == 'two')) self.client = httpx.Client(transport=httpx.MockTransport(respond)) self.addCleanup(self.client.close) def test_catalog_is_lazy_and_cache_is_bounded(self): source = GSOCloudSource(self.root, client=self.client, max_objects=1) self.assertEqual(len(source.objects), 2) self.assertEqual(len(self.requests), 1) self.assertEqual(list(source.root.glob('*.glb')), []) with source.open(source.objects[0]) as item: first = Path(item['mesh']) self.assertTrue(first.is_file()) metrics = cached_measure(item, self.root / 'metrics') self.assertAlmostEqual(metrics['a'], 3) with source.open(source.objects[0]): self.assertEqual(len(self.requests), 2) # Cache hit, no archive request. with source.open(source.objects[1]) as item: self.assertTrue(Path(item['mesh']).exists()) self.assertFalse(first.exists()) self.assertEqual(len(list(source.root.glob('*.glb'))), 1) with source.open(source.objects[0]) as item: with patch('labeler.dataset.measure', side_effect=AssertionError('Metrics must survive eviction')): self.assertEqual(cached_measure(item, self.root / 'metrics'), metrics) def test_failed_download_does_not_leave_a_partial_asset(self): source = GSOCloudSource(self.root, client=self.client) self.fail_archives = True with self.assertRaises(CloudDatasetError): with source.open(source.objects[0]): pass self.assertEqual(list(source.root.glob('*.glb')), []) self.assertEqual(list(source.root.glob('download-*')), []) self.fail_archives = False with source.open(source.objects[0]) as item: self.assertTrue(Path(item['mesh']).exists()) def test_cached_catalog_can_be_used_offline(self): online = GSOCloudSource(self.root, client=self.client) def offline(request): raise httpx.ConnectError('offline', request=request) with httpx.Client(transport=httpx.MockTransport(offline)) as client: with self.assertLogs(level='WARNING'): cached = GSOCloudSource(self.root, client=client) self.assertEqual(online.objects, cached.objects) with self.assertRaises(CloudDatasetError): GSOCloudSource(self.root / 'empty', client=client) def test_byte_limit_evicts_oversized_asset_after_use(self): source = GSOCloudSource(self.root, client=self.client, max_bytes=1) with source.open(source.objects[0]) as item: path = Path(item['mesh']) self.assertTrue(path.exists()) self.assertFalse(path.exists()) def test_archive_cannot_escape_cache_or_create_links(self): for name, kind in [('../escape', tarfile.REGTYPE), ('link', tarfile.SYMTYPE)]: archive_path = self.root / 'unsafe.tar.gz' with tarfile.open(archive_path, 'w:gz') as archive: member = tarfile.TarInfo(name) member.type = kind member.linkname = '/etc/passwd' if kind == tarfile.SYMTYPE else '' archive.addfile(member) with self.assertRaises(ValueError): unpack_archive(archive_path, self.root / 'extracted') self.assertFalse((self.root / 'escape').exists()) def test_default_app_uses_cloud_and_copies_viewer_before_eviction(self): os.environ['GRADIO_ANALYTICS_ENABLED'] = 'False' import gradio as gr from app import build_app from labeler.storage import AnnotationStore source = GSOCloudSource(self.root, client=self.client, max_bytes=1) annotations = self.root / 'annotations' with patch.dict(os.environ, {'DATA_SOURCE': '', 'DATA_DIR': '', 'DATA_REPO': '', 'ANNOTATION_REPO': '', 'ANNOTATOR_USERS': '{}', 'CACHE_DIR': str(self.root / 'metrics'), 'ANNOTATION_DIR': str(annotations)}): with patch('app.GSOCloudSource', return_value=source) as create: demo, _ = build_app() create.assert_called_once() self.assertEqual(len(self.requests), 1) # App startup fetched no archives. functions = {fn.fn.__name__: fn.fn for fn in demo.fns.values() if fn.fn} request = gr.Request() first = functions['render']('tester', 'Unlabeled', 'pca', 0, None, request) self.assertEqual(first[0], 'gso/one') self.assertTrue(Path(first[1]).is_file()) self.assertEqual(list(source.root.glob('*.glb')), []) second = functions['submit']('tester', 'Unlabeled', 'pca', 0, first[0], 'cuboid', 'sure', '', request) self.assertEqual(second[0], 'gso/two') self.assertTrue(Path(second[1]).is_file()) rows = AnnotationStore(annotations).records('tester') self.assertEqual(rows[0]['object_id'], 'gso/one') self.assertEqual(rows[0]['dataset_source'], 'gso')