aaronrene / scooling-lab public
test_unit.py python
276 lines 10.7 KB
Raw
sha256:ca1d0e687bff8686b37126eb8e4fb6d38b40e189352a4920d66ae89c7274e340 Add T4 job cancellation retry and validation Human minor ⚠ breaking 41 days ago
1 """Unit tier tests for Scooling Lab T0/T2 contracts."""
2
3 from __future__ import annotations
4
5 from datetime import UTC, datetime, timedelta
6 import unittest
7
8 from scooling_lab_helpers import valid_payload
9
10 from scooling_lab.contracts import TrainingJobRequest, TrainingJobStatus
11 from scooling_lab.dataset_review import (
12 RejectionReasonCode,
13 dataset_shape_from_registration,
14 validate_synthetic_dataset_shape,
15 )
16 from scooling_lab.errors import ApiError, ErrorCode, error_payload
17 from scooling_lab.license_policy import BomEntry, LicensePolicyError, validate_entry
18 from scooling_lab.provenance import ProvenanceRecord, validate_provenance_record
19 from scooling_lab.retention import (
20 RetentionPolicyClass,
21 expires_at,
22 is_expired,
23 retention_policy_from_mapping,
24 )
25 from scooling_lab.service import TrainingApiService
26 from scooling_lab.state_machine import cancel, transition
27 from scooling_lab.store import TrainingJobStore
28
29
30 class ScoolingLabUnitTests(unittest.TestCase):
31 """Unit tests for state transitions, schema, license policy, and errors."""
32
33 def test_unit_state_transitions_are_explicit_and_idempotent(self) -> None:
34 """queued -> running -> succeeded works and replay is a no-op."""
35
36 self.assertEqual(
37 transition(TrainingJobStatus.QUEUED, TrainingJobStatus.RUNNING),
38 TrainingJobStatus.RUNNING,
39 )
40 self.assertEqual(
41 transition(TrainingJobStatus.RUNNING, TrainingJobStatus.SUCCEEDED),
42 TrainingJobStatus.SUCCEEDED,
43 )
44 self.assertEqual(
45 transition(TrainingJobStatus.SUCCEEDED, TrainingJobStatus.SUCCEEDED),
46 TrainingJobStatus.SUCCEEDED,
47 )
48
49 def test_unit_invalid_state_transition_raises_safe_error(self) -> None:
50 """queued -> succeeded is not allowed by the transition table."""
51
52 with self.assertRaises(ApiError) as raised:
53 transition(TrainingJobStatus.QUEUED, TrainingJobStatus.SUCCEEDED)
54 self.assertEqual(raised.exception.code, ErrorCode.INVALID_TRANSITION)
55
56 def test_unit_t4_cancel_transitions_and_terminal_noops(self) -> None:
57 """Cancellation changes active states and leaves terminal states unchanged."""
58
59 self.assertEqual(cancel(TrainingJobStatus.QUEUED), TrainingJobStatus.CANCELLED)
60 self.assertEqual(cancel(TrainingJobStatus.RUNNING), TrainingJobStatus.CANCELLED)
61 for terminal in (
62 TrainingJobStatus.SUCCEEDED,
63 TrainingJobStatus.FAILED,
64 TrainingJobStatus.CANCELLED,
65 TrainingJobStatus.DELETED,
66 ):
67 with self.subTest(terminal=terminal):
68 self.assertEqual(cancel(terminal), terminal)
69
70 def test_unit_t4_retry_state_transitions_refuse_invalid_originals(self) -> None:
71 """Only failed or cancelled jobs can produce a fresh retry job."""
72
73 service = TrainingApiService(TrainingJobStore(), auto_run_worker=False)
74 queued = service.create_training_job(valid_payload("retry-invalid-queued"))
75 with self.assertRaises(ApiError) as queued_retry:
76 service.retry_training_job(str(queued["id"]))
77 self.assertEqual(queued_retry.exception.code, ErrorCode.INVALID_TRANSITION)
78
79 succeeded_service = TrainingApiService(TrainingJobStore())
80 succeeded = succeeded_service.create_training_job(valid_payload("retry-invalid-done"))
81 with self.assertRaises(ApiError) as succeeded_retry:
82 succeeded_service.retry_training_job(str(succeeded["id"]))
83 self.assertEqual(succeeded_retry.exception.code, ErrorCode.INVALID_TRANSITION)
84
85 def test_unit_t4_dataset_validation_maps_shapes_to_reason_codes(self) -> None:
86 """Synthetic dataset validation returns enum reasons for invalid shapes."""
87
88 _, valid_shape = dataset_shape_from_registration(
89 {
90 "datasetId": "unit-valid-shape",
91 "rowCount": 3,
92 "declaredSchema": {
93 "exampleId": "string",
94 "inputTokenCount": "integer",
95 "outputTokenCount": "integer",
96 "split": "string",
97 },
98 }
99 )
100 self.assertIsNone(validate_synthetic_dataset_shape(valid_shape))
101
102 _, too_large = dataset_shape_from_registration(
103 {
104 "datasetId": "unit-large-shape",
105 "rowCount": 10_001,
106 "declaredSchema": {
107 "exampleId": "string",
108 "inputTokenCount": "integer",
109 "outputTokenCount": "integer",
110 "split": "string",
111 },
112 }
113 )
114 self.assertEqual(
115 validate_synthetic_dataset_shape(too_large),
116 RejectionReasonCode.SYNTHETIC_LIMIT,
117 )
118
119 _, schema_mismatch = dataset_shape_from_registration(
120 {
121 "datasetId": "unit-schema-shape",
122 "rowCount": 3,
123 "declaredSchema": {
124 "exampleId": "string",
125 "inputTokenCount": "integer",
126 "split": "string",
127 },
128 }
129 )
130 self.assertEqual(
131 validate_synthetic_dataset_shape(schema_mismatch),
132 RejectionReasonCode.SCHEMA_MISMATCH,
133 )
134
135 def test_unit_deleted_state_transition_is_terminal(self) -> None:
136 """succeeded -> deleted works and deleted cannot transition elsewhere."""
137
138 self.assertEqual(
139 transition(TrainingJobStatus.SUCCEEDED, TrainingJobStatus.DELETED),
140 TrainingJobStatus.DELETED,
141 )
142 self.assertEqual(
143 transition(TrainingJobStatus.DELETED, TrainingJobStatus.DELETED),
144 TrainingJobStatus.DELETED,
145 )
146 with self.assertRaises(ApiError):
147 transition(TrainingJobStatus.DELETED, TrainingJobStatus.RUNNING)
148
149 def test_unit_schema_rejects_unsafe_fields_and_unapproved_models(self) -> None:
150 """Dangerous keys and unknown model ids fail in schema validation."""
151
152 payload = valid_payload()
153 payload["workerUrl"] = "https://attacker.invalid/worker"
154 with self.assertRaises(ApiError):
155 TrainingJobRequest.from_mapping(payload)
156
157 model_payload = valid_payload()
158 model_payload["modelId"] = "unapproved-model"
159 with self.assertRaises(ApiError):
160 TrainingJobRequest.from_mapping(model_payload)
161
162 def test_unit_provenance_schema_rejects_free_text_paths_and_urls(self) -> None:
163 """Provenance accepts only exact hashes, ids, timestamps, and schema version."""
164
165 record = ProvenanceRecord(
166 artifact_hash="b" * 64,
167 base_model_id="fixture-tiny-llm",
168 created_at="2026-06-11T00:00:00Z",
169 dataset_hash="a" * 64,
170 job_id="job_" + "1" * 24,
171 training_config_hash="c" * 64,
172 )
173 validate_provenance_record(record)
174
175 unsafe_values = (
176 "prompt body with words",
177 "../private/path",
178 "https://attacker.invalid/model",
179 )
180 for unsafe in unsafe_values:
181 payload = record.to_dict()
182 payload["baseModelId"] = unsafe
183 with self.subTest(unsafe=unsafe):
184 with self.assertRaises(ApiError):
185 validate_provenance_record(payload)
186
187 def test_unit_retention_ttl_math_is_bounded_and_deterministic(self) -> None:
188 """Retention policies validate TTL bounds and evaluate expiry exactly."""
189
190 policy = retention_policy_from_mapping(
191 {"policyClass": "ephemeral", "ttlSeconds": 60}
192 )
193 self.assertEqual(policy.policy_class, RetentionPolicyClass.EPHEMERAL)
194 self.assertEqual(expires_at("2026-06-11T00:00:00Z", policy), "2026-06-11T00:01:00Z")
195 self.assertFalse(
196 is_expired(
197 "2026-06-11T00:00:00Z",
198 policy,
199 datetime(2026, 6, 11, 0, 0, 59, tzinfo=UTC),
200 )
201 )
202 self.assertTrue(
203 is_expired(
204 "2026-06-11T00:00:00Z",
205 policy,
206 datetime(2026, 6, 11, 0, 1, 0, tzinfo=UTC),
207 )
208 )
209 with self.assertRaises(ApiError):
210 retention_policy_from_mapping({"policyClass": "ephemeral", "ttlSeconds": 59})
211
212 def test_unit_deletion_state_transition_removes_content_fields(self) -> None:
213 """deleteArtifact leaves a content-free deleted job tombstone."""
214
215 service = TrainingApiService(TrainingJobStore())
216 created = service.create_training_job(valid_payload("delete-unit"))
217 job_id = str(created["id"])
218 artifact_id = str(service.list_artifacts(job_id)["artifacts"][0]["id"])
219
220 first_delete = service.delete_artifact(job_id, artifact_id)
221 second_delete = service.delete_artifact(job_id, artifact_id)
222 tombstone = service.get_training_job(job_id)
223
224 self.assertTrue(first_delete["deleted"])
225 self.assertTrue(first_delete["verified"])
226 self.assertTrue(second_delete["alreadyDeleted"])
227 self.assertEqual(tombstone["status"], "deleted")
228 self.assertNotIn("request", tombstone)
229 self.assertNotIn("artifactHash", str(tombstone))
230
231 def test_unit_license_policy_blocks_agpl_and_blocked_paths(self) -> None:
232 """The BOM policy rejects AGPL licenses and Unsloth Studio/CLI paths."""
233
234 validate_entry(
235 BomEntry(
236 name="safe",
237 version="1.0.0",
238 license="MIT",
239 source_path="third_party/safe",
240 evidence="fixture",
241 )
242 )
243 with self.assertRaises(LicensePolicyError):
244 validate_entry(
245 BomEntry(
246 name="blocked",
247 version="1.0.0",
248 license="AGPL-3.0",
249 source_path="third_party/blocked",
250 evidence="fixture",
251 )
252 )
253 with self.assertRaises(LicensePolicyError):
254 validate_entry(
255 BomEntry(
256 name="blocked-path",
257 version="1.0.0",
258 license="Apache-2.0",
259 source_path="vendor/unsloth_cli/main.py",
260 evidence="fixture",
261 )
262 )
263
264 def test_unit_error_model_does_not_echo_payload_or_paths(self) -> None:
265 """Safe errors expose only code and stable public message."""
266
267 payload = error_payload(ApiError(ErrorCode.VALIDATION_ERROR, 400))
268 text = str(payload)
269 self.assertIn("VALIDATION_ERROR", text)
270 self.assertNotIn("/Users/", text)
271 self.assertNotIn("attacker.invalid", text)
272 self.assertNotIn("payload", text.lower())
273
274
275 if __name__ == "__main__":
276 unittest.main()
File History 1 commit
sha256:ca1d0e687bff8686b37126eb8e4fb6d38b40e189352a4920d66ae89c7274e340 Add T4 job cancellation retry and validation Human minor ⚠ 41 days ago