IntReal-Leaderboard / tests /test_submissions.py
Sean Yu
Limit each HF account to three submissions per seven days
433597d
Raw History Blame Contribute Delete
14 kB
import asyncio
import base64
import io
import json
import stat
import time
import zipfile
import httpx
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from itsdangerous import TimestampSigner
from PIL import Image
import submission_backend as backend
def archive_bytes(name='images/0.png'):
image = io.BytesIO()
Image.new('RGB', (8, 8), 'blue').save(image, format='PNG')
data = io.BytesIO()
with zipfile.ZipFile(data, 'w') as archive:
archive.writestr(name, image.getvalue())
archive.writestr('outputs.jsonl', '{"id": 1}\n')
return data.getvalue()
def test_model_owner_and_url():
assert backend.model_id_from_url('https://huggingface.co/alice/model', 'alice') == 'alice/model'
for value in ['https://huggingface.co/bob/model', 'https://evil.test/alice/model', 'https://huggingface.co@evil.test/alice/model', 'https://huggingface.co/datasets/alice/model', 'https://huggingface.co/alice/model?x=1']:
with pytest.raises(ValueError):
backend.model_id_from_url(value, 'alice')
@pytest.mark.parametrize('changes', [{'private': True}, {'gated': 'auto'}, {'cardData': {'license': 'other'}}, {'siblings': []}])
def test_model_rejects_private_gated_nonopen_or_missing_weights(changes):
info = {'id': 'alice/model', 'private': False, 'gated': False, 'cardData': {'license': 'apache-2.0'}, 'siblings': [{'rfilename': 'model.safetensors'}], **changes}
async def run():
async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200, json=info))) as client:
with pytest.raises(ValueError):
await backend.validate_model('alice/model', client)
asyncio.run(run())
def test_model_check_is_anonymous():
def respond(request):
assert request.headers.get('authorization') == ''
return httpx.Response(200, json={'id': 'alice/model', 'private': False, 'gated': False, 'sha': 'abc', 'cardData': {'license': 'mit'}, 'siblings': [{'rfilename': 'model.safetensors'}]})
async def run():
async with httpx.AsyncClient(transport=httpx.MockTransport(respond), headers={'Authorization': 'Bearer must-not-be-used'}) as client:
assert (await backend.validate_model('alice/model', client))['revision'] == 'abc'
asyncio.run(run())
def test_valid_zip(tmp_path):
path = tmp_path / 'test.zip'
path.write_bytes(archive_bytes())
assert backend.inspect_zip(path)['image_count'] == 1
@pytest.mark.parametrize('name', ['../bad.png', '/bad.png', 'C:/bad.png', 'foo\\bad.png', 'run.py'])
def test_unsafe_zip_entries(tmp_path, name):
path = tmp_path / 'test.zip'
path.write_bytes(archive_bytes(name))
with pytest.raises(ValueError):
backend.inspect_zip(path)
def test_invalid_and_empty_zip(tmp_path):
path = tmp_path / 'test.zip'
path.write_text('not really zip')
with pytest.raises(ValueError):
backend.inspect_zip(path)
with zipfile.ZipFile(path, 'w') as archive:
archive.writestr('metadata.json', '{}')
with pytest.raises(ValueError, match='none were found'):
backend.inspect_zip(path)
def test_fake_image_and_symlink(tmp_path):
path = tmp_path / 'test.zip'
with zipfile.ZipFile(path, 'w') as archive:
archive.writestr('fake.png', b'not image')
with pytest.raises(ValueError):
backend.inspect_zip(path)
with zipfile.ZipFile(path, 'w') as archive:
item = zipfile.ZipInfo('link.png')
item.create_system = 3
item.external_attr = (stat.S_IFLNK | 0o777) << 16
archive.writestr(item, 'target')
with pytest.raises(ValueError):
backend.inspect_zip(path)
def cookie(username='alice', csrf='csrf-test'):
value = {'hf_user': {'username': username, 'sub': 'hf-'+username, 'expires_at': time.time()+3600}, 'csrf': csrf}
raw = base64.b64encode(json.dumps(value).encode())
return TimestampSigner('test-secret').sign(raw).decode()
@pytest.fixture
def service(tmp_path, monkeypatch):
monkeypatch.setenv('SESSION_SECRET', 'test-secret')
monkeypatch.setenv('SUBMISSIONS_DIR', str(tmp_path / 'private'))
monkeypatch.setenv('INTREAL_LOCAL_STORAGE', '1')
monkeypatch.delenv('SPACE_ID', raising=False)
monkeypatch.delenv('SPACE_HOST', raising=False)
app = FastAPI()
backend.install_submissions(app)
async def verified(repo_id, client):
return {'repo_id': repo_id, 'revision': 'abc', 'license': 'mit'}
monkeypatch.setattr(backend, 'validate_model', verified)
with TestClient(app) as client:
yield client, tmp_path / 'private'
def form(data):
return {'model_url': 'https://huggingface.co/alice/model', 'username': 'alice', 'email': 'alice@example.com', 'filename': 'outputs.zip', 'size': len(data)}
def sign_in(client, username='alice'):
client.cookies.clear()
client.cookies.set('intreal_session', cookie(username))
def test_auth_and_csrf_required(service):
client, _ = service
assert client.post('/api/submissions', json=form(archive_bytes())).status_code == 401
sign_in(client)
assert client.post('/api/submissions', json=form(archive_bytes())).status_code == 403
assert client.put('/api/submissions/'+'a'*32+'/archive', content=b'zip').status_code == 403
def test_forged_username_rejected(service):
client, _ = service
sign_in(client)
payload = form(archive_bytes()) | {'username': 'bob'}
response = client.post('/api/submissions', json=payload, headers={'X-CSRF-Token':'csrf-test'})
assert response.status_code == 422
def test_receive_private_archive_and_no_public_download(service):
client, root = service
sign_in(client)
data = archive_bytes()
headers = {'X-CSRF-Token': 'csrf-test'}
prepared = client.post('/api/submissions', json=form(data), headers=headers)
assert prepared.status_code == 200, prepared.text
sid = prepared.json()['id']
sign_in(client, 'bob')
assert client.put(f'/api/submissions/{sid}/archive', content=data, headers=headers).status_code == 404
sign_in(client)
received = client.put(f'/api/submissions/{sid}/archive', content=data, headers=headers)
assert received.status_code == 200, received.text
assert received.json()['status'] == 'pending_review'
assert (root / sid / 'outputs.zip').read_bytes() == data
record = json.loads((root / sid / 'metadata.json').read_text())
assert record['username'] == 'alice' and record['archive']['image_count'] == 1
assert not (root / sid / 'images').exists()
assert client.get(f'/api/submissions/{sid}').status_code == 405
assert client.get(f'/leaderboard-assets/{sid}/outputs.zip').status_code == 404
assert client.put(f'/api/submissions/{sid}/archive', content=data, headers=headers).status_code == 409
def test_bad_archive_cleanup(service):
client, root = service
sign_in(client)
data = b'fake zip'
headers = {'X-CSRF-Token': 'csrf-test'}
sid = client.post('/api/submissions', json=form(data), headers=headers).json()['id']
assert client.put(f'/api/submissions/{sid}/archive', content=data, headers=headers).status_code == 422
assert not list((root / sid).glob('*.part'))
assert not (root / sid / 'outputs.zip').exists()
assert json.loads((root / sid / 'metadata.json').read_text())['status'] == 'failed'
def test_truncated_upload_and_size_limit(service):
client, root = service
sign_in(client)
data = archive_bytes()
headers = {'X-CSRF-Token': 'csrf-test'}
too_large = form(data) | {'size': backend.MAX_ZIP + 1}
assert client.post('/api/submissions', json=too_large, headers=headers).status_code == 422
sid = client.post('/api/submissions', json=form(data), headers=headers).json()['id']
assert client.put(f'/api/submissions/{sid}/archive', content=data[:-10], headers=headers).status_code == 422
assert not (root / sid / 'outputs.zip').exists()
def test_no_persistence_fails_closed(tmp_path, monkeypatch):
monkeypatch.setenv('SPACE_ID', 'owner/space')
monkeypatch.setenv('INTREAL_LOCAL_STORAGE', '1')
monkeypatch.setenv('SUBMISSIONS_DIR', str(tmp_path))
monkeypatch.setenv('SUBMISSIONS_MOUNT', str(tmp_path / 'not-a-mount'))
assert backend.SubmissionService().ready() is False
def test_forged_cookie_is_not_identity(service):
client, _ = service
client.cookies.set('intreal_session', cookie()[:-8] + 'forged!!')
assert client.post('/api/submissions', json=form(archive_bytes()), headers={'X-CSRF-Token': 'csrf-test'}).status_code == 401
def test_duplicate_paths_and_expansion_limit(tmp_path, monkeypatch):
path = tmp_path / 'test.zip'
with zipfile.ZipFile(path, 'w') as archive:
archive.writestr('same.txt', 'a')
with pytest.warns(UserWarning):
archive.writestr('same.txt', 'b')
with pytest.raises(ValueError, match='duplicate'):
backend.inspect_zip(path)
path.write_bytes(archive_bytes())
monkeypatch.setattr(backend, 'MAX_EXPANDED', 1)
with pytest.raises(ValueError, match='expanded-size'):
backend.inspect_zip(path)
def test_single_pending_and_weekly_limit(service):
client, _ = service
sign_in(client)
headers = {'X-CSRF-Token': 'csrf-test'}
data = archive_bytes()
sid = client.post('/api/submissions', json=form(data), headers=headers).json()['id']
assert client.post('/api/submissions', json=form(data), headers=headers).status_code == 409
assert client.delete('/api/submissions/'+sid, headers=headers).status_code == 200
for _ in range(3):
sid = client.post('/api/submissions', json=form(b'invalid'), headers=headers).json()['id']
assert client.put(f'/api/submissions/{sid}/archive', content=b'invalid', headers=headers).status_code == 422
assert client.post('/api/submissions', json=form(data), headers=headers).status_code == 429
def test_server_quota_includes_reserved_uploads(service, monkeypatch):
client, _ = service
sign_in(client)
data = archive_bytes()
monkeypatch.setattr(backend, 'MAX_TOTAL', len(data)-1)
response = client.post('/api/submissions', json=form(data), headers={'X-CSRF-Token': 'csrf-test'})
assert response.status_code == 507
def test_invalid_email_and_non_zip_name(service):
client, _ = service
sign_in(client)
for patch in [{'email': 'not-email'}, {'filename': 'outputs.tar'}, {'model_url': 'https://huggingface.co/bob/model'}]:
response = client.post('/api/submissions', json=form(archive_bytes()) | patch, headers={'X-CSRF-Token':'csrf-test'})
assert response.status_code == 422
def test_login_disabled_without_real_oauth(service):
client, _ = service
assert client.get('/api/submission-session').json()['auth_enabled'] is False
assert client.get('/auth/login').status_code == 503
def test_full_storage_message_and_disabled_state(service, monkeypatch):
client, _ = service
sign_in(client)
monkeypatch.setattr(backend, 'MAX_TOTAL', 0)
state = client.get('/api/submission-session').json()
assert state['storage_full'] is True
assert 'shaohanyu628@gmail.com' in state['storage_message']
response = client.post('/api/submissions', json=form(archive_bytes()), headers={'X-CSRF-Token':'csrf-test'})
assert response.status_code == 507
assert 'submit again' in response.json()['detail']
def test_low_disk_blocks_upload_before_bytes(service, monkeypatch):
from collections import namedtuple
client, _ = service
sign_in(client)
disk = namedtuple('usage', 'total used free')
monkeypatch.setattr(backend.shutil, 'disk_usage', lambda _: disk(10, 9, 1))
response = client.post('/api/submissions', json=form(archive_bytes()), headers={'X-CSRF-Token':'csrf-test'})
assert response.status_code == 507
assert 'shaohanyu628@gmail.com' in response.json()['detail']
def test_disk_fills_during_upload_cleans_partial(service, monkeypatch):
from collections import namedtuple
client, root = service
sign_in(client)
data = archive_bytes()
headers = {'X-CSRF-Token':'csrf-test'}
sid = client.post('/api/submissions', json=form(data), headers=headers).json()['id']
disk = namedtuple('usage', 'total used free')
monkeypatch.setattr(backend.shutil, 'disk_usage', lambda _: disk(10, 9, 1))
response = client.put(f'/api/submissions/{sid}/archive',content=data,headers=headers)
assert response.status_code == 507
assert 'shaohanyu628@gmail.com' in response.json()['detail']
assert not (root/sid/'outputs.zip.part').exists()
assert not (root/sid/'outputs.zip').exists()
def test_model_name_is_derived_and_not_user_controlled(service):
client, root = service
sign_in(client)
payload = form(archive_bytes()) | {'model_name': 'forged display name'}
response = client.post('/api/submissions', json=payload, headers={'X-CSRF-Token':'csrf-test'})
assert response.status_code == 200
record = json.loads((root / response.json()['id'] / 'metadata.json').read_text())
assert record['model_name'] == 'model'
def test_submission_quota_uses_rolling_seven_days(service):
client, root = service
sign_in(client)
headers = {'X-CSRF-Token':'csrf-test'}
data = archive_bytes()
for index in range(3):
sid = f'{index:032x}'
folder = root / sid
folder.mkdir(parents=True)
(folder / 'metadata.json').write_text(json.dumps({
'id': sid, 'username': 'alice', 'created_at': time.time() - 6*86400,
'status': 'failed', 'declared_bytes': 10,
}))
response = client.post('/api/submissions', json=form(data), headers=headers)
assert response.status_code == 429
assert '7-day' in response.json()['detail']
for path in root.glob('*/metadata.json'):
record = json.loads(path.read_text())
record['created_at'] = time.time() - 8*86400
path.write_text(json.dumps(record))
assert client.post('/api/submissions', json=form(data), headers=headers).status_code == 200