| 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) |
| 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) |
| 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') |
|
|
