"""Private disk layout, 50-Go private cap, and real-byte pooled accounting. Run with: USE_SQLITE=true python manage.py test drive.test_private_layout """ import tempfile import uuid from pathlib import Path from django.contrib.auth import get_user_model from django.core.files.uploadedfile import SimpleUploadedFile from django.core.management import call_command from django.core.management.base import CommandError from django.db import transaction from django.test import TestCase, override_settings from rest_framework.test import APIClient from .models import DriveItem, DriveSpace, SpaceInvitation, SpaceGrant from .quotas import (PRIVATE_SPACE_QUOTA_BYTES, QuotaExceeded, assert_space_capacity, assert_updated_quota, locked_quota_tree, occupancy, quota_state) from .spaces import ensure_private_space, ensure_space, create_space, storage_path class PrivateStorageTests(TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.addCleanup(self.temp.cleanup) self.rootdir = Path(self.temp.name) overrides = override_settings( DRIVE_STORAGE_ROOT=self.rootdir, FILE_UPLOAD_TEMP_DIR=str(self.rootdir / '.uploads'), ) overrides.enable() self.addCleanup(overrides.disable) User = get_user_model() self.admin = User.objects.create_superuser('root-admin', 'admin@example.org', 'TestPassword123!') self.shared = ensure_space(self.admin) self.private = ensure_private_space(self.admin) self.client = APIClient() self.client.force_authenticate(user=self.admin) def test_private_directories_are_separate_from_shared_spaces(self): self.assertEqual(self.private.parent_id, self.shared.pk) self.assertEqual(self.private.relative_path, f'privates/{self.private.directory_name}') self.assertEqual(self.private.quota_bytes, PRIVATE_SPACE_QUOTA_BYTES) self.assertTrue((self.rootdir / self.private.relative_path).is_dir()) person = get_user_model().objects.create_user('member', password='TestPassword123!') shared = create_space(person, self.shared, 'Compte membre') private = ensure_private_space(person) self.assertTrue(shared.relative_path.startswith(f'{shared.directory_name}')) self.assertEqual(private.relative_path, f'privates/{private.directory_name}') self.assertEqual(private.parent_id, shared.pk) self.assertNotEqual(shared.directory_name, private.directory_name) def test_private_usage_counts_toward_assigned_space_without_reserving(self): self.shared.quota_bytes = 100 self.shared.save(update_fields=['quota_bytes']) # Private can be 50 Go even when the assigned budget is just 100 bytes. self.assertEqual(quota_state(self.shared)['available_bytes'], 100) self.assertEqual(quota_state(self.private)['available_bytes'], 100) created = self.client.post('/api/drive/uploads/', { 'files': [SimpleUploadedFile('photo.jpg', b'abcdef')] }, format='multipart') self.assertEqual(created.status_code, 201, created.data) node = DriveItem.objects.get(pk=created.data['items'][0]['id']) self.assertEqual(node.storage_location, f'{self.private.relative_path}/photo.jpg') self.assertEqual(storage_path(node).read_bytes(), b'abcdef') self.assertEqual(occupancy(self.shared)[0], 6) self.assertEqual(quota_state(self.shared)['available_bytes'], 94) self.assertEqual(quota_state(self.private)['available_bytes'], 94) self.assertEqual(quota_state(self.private)['quota_bytes'], PRIVATE_SPACE_QUOTA_BYTES) self.assertEqual(quota_state(self.private)['used_bytes'], 6) rejected = self.client.patch(f'/api/drive/spaces/{self.private.pk}/quota/', { 'quota_bytes': None }, format='json') self.assertEqual(rejected.status_code, 403) def test_actual_private_use_is_pooled_in_100_go_assigned_space(self): person = get_user_model().objects.create_user('quota-user', password='TestPassword123!') assigned = create_space(person, self.shared, '100 Go', quota_bytes=100 * 1024 ** 3) private = ensure_private_space(person) gib = 1024 ** 3 self.assertEqual(private.parent_id, assigned.pk) self.assertEqual(private.relative_path, f'privates/{private.directory_name}') self.assertEqual(quota_state(assigned)['available_bytes'], 100 * gib) self.assertEqual(quota_state(private)['available_bytes'], 50 * gib) # Existing 20 Go in shared space. This leaves both limits at their max. DriveItem.objects.create( owner=person, space=assigned, name='shared.bin', storage_key=str(uuid.uuid4()), storage_location=f'{assigned.relative_path}/shared.bin', size=20 * gib, ) self.assertEqual(quota_state(assigned)['used_bytes'], 20 * gib) self.assertEqual(quota_state(private)['used_bytes'], 0) DriveItem.objects.create( owner=person, space=private, name='private.bin', storage_key=str(uuid.uuid4()), storage_location=f'{private.relative_path}/private.bin', size=20 * gib, ) self.assertEqual(occupancy(assigned)[0], 40 * gib) self.assertEqual(quota_state(assigned)['used_bytes'], 40 * gib) self.assertEqual(quota_state(assigned)['available_bytes'], 60 * gib) self.assertEqual(quota_state(private)['used_bytes'], 20 * gib) self.assertEqual(quota_state(private)['available_bytes'], 30 * gib) with transaction.atomic(): with locked_quota_tree(assigned): assert_space_capacity(assigned, 60 * gib) with self.assertRaises(QuotaExceeded): assert_space_capacity(assigned, 60 * gib + 1) assert_updated_quota(assigned, 40 * gib) with self.assertRaisesRegex(ValueError, 'fichiers'): assert_updated_quota(assigned, 39 * gib) def test_parent_limit_blocks_private_even_before_50_go(self): self.shared.quota_bytes = 20 self.shared.save(update_fields=['quota_bytes']) DriveItem.objects.create( owner=self.admin, space=self.private, name='existing.bin', storage_key=str(uuid.uuid4()), storage_location=f'{self.private.relative_path}/existing.bin', size=19, ) self.assertEqual(quota_state(self.private)['available_bytes'], 1) with transaction.atomic(): with locked_quota_tree(self.private): assert_space_capacity(self.private, 1) with self.assertRaises(QuotaExceeded): assert_space_capacity(self.private, 2) def test_invitation_preview_includes_creator_full_name_and_parent_space(self): self.admin.first_name = 'Camille' self.admin.last_name = 'Martin' self.admin.save(update_fields=['first_name', 'last_name']) result = self.client.post('/api/drive/invitations/', { 'parent_space': self.shared.pk, }, format='json') self.assertEqual(result.status_code, 201, result.data) preview = APIClient().get(f"/api/invitations/{result.data['token']}/") self.assertEqual(preview.status_code, 200, preview.data) self.assertEqual(preview.data['inviter_name'], 'Camille Martin') self.assertEqual(preview.data['parent_name'], self.shared.display_name) self.admin.first_name = '' self.admin.last_name = '' self.admin.save(update_fields=['first_name', 'last_name']) preview = APIClient().get(f"/api/invitations/{result.data['token']}/") self.assertEqual(preview.status_code, 200, preview.data) self.assertEqual(preview.data['inviter_name'], self.admin.username) def test_invitation_has_no_quota_even_from_an_old_client(self): self.shared.quota_bytes = 15 self.shared.save(update_fields=['quota_bytes']) # Even an old client cannot allocate storage through an invitation. result = self.client.post('/api/drive/invitations/', { 'parent_space': self.shared.pk, 'quota_bytes': 10, }, format='json') self.assertEqual(result.status_code, 201, result.data) self.assertNotIn('quota_bytes', result.data) invitation = SpaceInvitation.objects.get(pk=result.data['id']) self.assertIsNone(invitation.quota_bytes) self.assertEqual(quota_state(self.shared)['available_bytes'], 15) # Legacy unredeemed invitations may still contain a saved quota. # The quota must not be applied when the invite is accepted. invitation.quota_bytes = 1 invitation.save(update_fields=['quota_bytes']) preview = APIClient().get(f"/api/invitations/{result.data['token']}/") self.assertEqual(preview.status_code, 200, preview.data) self.assertNotIn('quota_bytes', preview.data) accepted = APIClient().post( f"/api/invitations/{result.data['token']}/accept/", { 'username': 'invitee', 'password': 'TestPassword123!', 'first_name': 'Jean', 'last_name': 'Dupont', 'display_name': 'Dossier invité', }, format='json', ) self.assertEqual(accepted.status_code, 201, accepted.data) user = get_user_model().objects.get(username='invitee') self.assertFalse(DriveSpace.objects.filter(user=user).exists()) self.assertTrue(SpaceGrant.objects.filter(user=user, space=self.shared).exists()) assigned = ensure_space(user) private = ensure_private_space(user) self.assertEqual(assigned.pk, self.shared.pk) self.assertEqual(private.quota_bytes, PRIVATE_SPACE_QUOTA_BYTES) self.assertEqual(private.parent_id, self.shared.pk) self.assertEqual(quota_state(assigned)['available_bytes'], 15) self.assertEqual(quota_state(private)['available_bytes'], 15) def test_private_pending_uploads_count_in_parent(self): from .models import UploadSession UploadSession.objects.create( owner=self.admin, space=self.private, name='movie.bin', total_size=1234 ) self.assertEqual(quota_state(self.shared)['reserved_bytes'], 1234) self.assertEqual(quota_state(self.shared)['used_bytes'], 0) self.assertEqual(quota_state(self.private)['available_bytes'], PRIVATE_SPACE_QUOTA_BYTES - 1234) def test_private_50_go_upload_limit_is_enforced(self): DriveItem.objects.create( owner=self.admin, space=self.private, name='large.bin', storage_key=str(uuid.uuid4()), storage_location=f'{self.private.relative_path}/large.bin', size=PRIVATE_SPACE_QUOTA_BYTES - 1, ) with transaction.atomic(): with locked_quota_tree(self.private): with self.assertRaises(QuotaExceeded): assert_space_capacity(self.private, 2) assert_space_capacity(self.private, 1) def simulate_legacy_private(self): """Relocate the empty modern directory to the historical nested location.""" new = self.rootdir / self.private.relative_path new.rmdir() old_name = self.private.directory_name old = f'{old_name}' src = self.rootdir / old (src / 'Documents').mkdir(parents=True) (src / 'Documents' / 'facture.pdf').write_bytes(b'PDF CONTENT') self.private.parent = self.shared self.private.relative_path = old self.private.quota_bytes = None self.private.save(update_fields=['parent', 'relative_path', 'quota_bytes']) folder = DriveItem.objects.create( owner=self.admin, space=self.private, name='Documents', is_folder=True, ) item = DriveItem.objects.create( owner=self.admin, space=self.private, parent=folder, name='facture.pdf', storage_key=str(uuid.uuid4()), storage_location=f'{old}/Documents/facture.pdf', size=11, ) return old, item def test_migrate_private_storage_dry_run_and_apply(self): old, item = self.simulate_legacy_private() call_command('move_private_storage') self.private.refresh_from_db() self.assertEqual(self.private.relative_path, old) self.assertTrue((self.rootdir / old / 'Documents' / 'facture.pdf').exists()) call_command('move_private_storage', '--apply') self.private.refresh_from_db() item.refresh_from_db() self.assertEqual(self.private.relative_path, f'privates/{self.private.directory_name}') self.assertEqual(self.private.parent_id, self.shared.pk) self.assertEqual(self.private.quota_bytes, PRIVATE_SPACE_QUOTA_BYTES) self.assertEqual(item.storage_location, f'{self.private.relative_path}/Documents/facture.pdf') self.assertEqual(storage_path(item).read_bytes(), b'PDF CONTENT') self.assertFalse((self.rootdir / old).exists()) self.assertFalse((self.rootdir / '.gooddrive-move-privates.json').exists()) call_command('move_private_storage', '--apply') # Idempotent. def test_interrupted_migration_can_resume_after_directory_move(self): from .management.commands.move_private_storage import Command, write_journal old, item = self.simulate_legacy_private() plan = Command().plan() journal = self.rootdir / '.gooddrive-move-privates.json' write_journal(journal, plan) new = self.rootdir / 'privates' / self.private.directory_name (self.rootdir / old).rename(new) # Simulate a crash before SQL update. call_command('move_private_storage', '--apply') item.refresh_from_db() self.assertEqual(storage_path(item).read_bytes(), b'PDF CONTENT') self.assertFalse(journal.exists()) def test_migration_never_overwrites_existing_destination(self): old, item = self.simulate_legacy_private() target = self.rootdir / 'privates' / self.private.directory_name target.mkdir(parents=True) (target / 'unrelated.txt').write_text('KEEP ME') with self.assertRaises(CommandError): call_command('move_private_storage', '--apply') self.assertEqual((target / 'unrelated.txt').read_text(), 'KEEP ME') self.assertTrue((self.rootdir / old / 'Documents' / 'facture.pdf').exists()) item.refresh_from_db() self.assertEqual(item.storage_location, f'{old}/Documents/facture.pdf')