Spaces:
Sleeping
Sleeping
Cline
Usability + abuse protection: run limits, rate limits, login by username/email, authed artifacts, presets, dataset search, stale-run cleanup, tests
1379106 | import os | |
| from aitrainlab import config | |
| from conftest import real_run | |
| CSV = ("text,label\n" + "\n".join(f"hello {i},{i % 2}" for i in range(30))).encode() | |
| def test_artifacts_need_login_and_ownership(client, auth): | |
| a, b = auth("alice"), auth("bob") | |
| ds = client.post("/datasets", json={"name": "mine"}, headers=a).json() | |
| url = f"/artifacts/dataset/{ds['id']}" | |
| assert client.post(url, files={"file": ("d.csv", CSV)}).status_code == 401 | |
| assert client.post(url, files={"file": ("d.csv", CSV)}, headers=b).status_code == 403 | |
| up = client.post(url, files={"file": ("d.csv", CSV)}, headers=a) | |
| assert up.status_code == 201 and up.json()["size_bytes"] == len(CSV) | |
| assert client.get(url).status_code == 401 | |
| assert client.get(url + "/d.csv").status_code == 401 | |
| assert client.get(url + "/d.csv", headers=b).status_code == 403 | |
| assert client.get(url + "/d.csv", headers=a).content == CSV | |
| assert client.post(url, files={"file": ("evil.exe", b"x")}, headers=a).status_code == 415 | |
| assert client.post(url, files={"file": ("../../x.csv", CSV)}, headers=a).status_code == 201 # basename only | |
| assert client.delete(url + "/d.csv", headers=a).status_code == 204 | |
| assert client.get(url + "/d.csv", headers=a).status_code == 404 | |
| def test_upload_size_limit(client, auth, monkeypatch): | |
| monkeypatch.setattr(config, "MAX_UPLOAD_MB", 1) | |
| h = auth() | |
| ds = client.post("/datasets", json={"name": "big"}, headers=h).json() | |
| r = client.post(f"/artifacts/dataset/{ds['id']}", files={"file": ("big.csv", b"a" * (2 * 1024 * 1024))}, headers=h) | |
| assert r.status_code == 413 | |
| assert not os.path.exists(os.path.join(config.STORAGE_DIR, "dataset", ds["id"], "big.csv")) | |
| def test_uploaded_dataset_run_requires_own_existing_file(client, auth): | |
| a, b = auth("alice"), auth("bob") | |
| ds = client.post("/datasets", json={"name": "mine"}, headers=a).json() | |
| client.post(f"/artifacts/dataset/{ds['id']}", files={"file": ("d.csv", CSV)}, headers=a) | |
| ref = f"upload:{ds['id']}/d.csv" | |
| assert client.post("/runs", json=real_run(dataset=ref), headers=b).status_code == 422 # not bob's | |
| assert client.post("/runs", json=real_run(dataset=f"upload:{ds['id']}/missing.csv"), headers=a).status_code == 422 | |
| assert client.post("/runs", json=real_run(dataset=ref), headers=a).status_code == 201 | |
| def test_upstream_requires_login(client): | |
| assert client.get("/upstream/search?q=bert").status_code == 401 | |
| assert client.get("/upstream/search-datasets?q=bert").status_code == 401 | |
| assert client.post("/upstream/import", json={"repo_id": "x/y"}).status_code == 401 | |
| def test_local_dataset_resolver_blocks_traversal(monkeypatch): | |
| import sys | |
| import types | |
| # transformers is heavy and irrelevant here; stub it so the module imports in CI. | |
| fake = types.ModuleType("transformers") | |
| fake.TrainerCallback = object | |
| monkeypatch.setitem(sys.modules, "transformers", fake) | |
| from aitrainlab.training.transformers_backend import _resolve_local | |
| assert _resolve_local("/etc/passwd") is None | |
| assert _resolve_local("upload:../x/passwd") is None | |
| assert _resolve_local("upload:abc/../../secret.csv") is None | |