123 lines
4.8 KiB
Python
123 lines
4.8 KiB
Python
import hashlib
|
|
import io
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
import controller
|
|
|
|
|
|
class ControllerTests(unittest.TestCase):
|
|
def test_get_requests_encode_parameters_in_the_query_string(self):
|
|
response = io.BytesIO(json.dumps({"code": "Success", "data": "stopped"}).encode())
|
|
client = controller.AutoDLClient("test-token")
|
|
with mock.patch.object(
|
|
controller.urllib.request, "urlopen", return_value=response
|
|
) as urlopen:
|
|
self.assertEqual(client.status("pro-test"), "stopped")
|
|
|
|
request = urlopen.call_args.args[0]
|
|
self.assertEqual(request.method, "GET")
|
|
self.assertIsNone(request.data)
|
|
self.assertTrue(request.full_url.endswith("?instance_uuid=pro-test"))
|
|
|
|
def test_runner_download_uses_explicit_user_agent(self):
|
|
payload = b"test-runner"
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
paths = controller.Paths(
|
|
root / "model",
|
|
"/model",
|
|
"/artifacts",
|
|
root / "artifacts",
|
|
root / "key",
|
|
root / "known_hosts",
|
|
root / "state.json",
|
|
root / "runner",
|
|
)
|
|
with mock.patch.object(
|
|
controller, "RUNNER_SHA256", hashlib.sha256(payload).hexdigest()
|
|
), mock.patch.object(
|
|
controller.urllib.request, "urlopen", return_value=io.BytesIO(payload)
|
|
) as urlopen:
|
|
controller.ensure_runner(paths)
|
|
|
|
request = urlopen.call_args.args[0]
|
|
self.assertEqual(request.get_header("User-agent"), controller.RUNNER_USER_AGENT)
|
|
|
|
def test_list_images_accepts_null_for_an_empty_account(self):
|
|
client = controller.AutoDLClient("test-token")
|
|
client.request = lambda method, path, body: {"list": None}
|
|
|
|
self.assertEqual(client.list_images(), [])
|
|
|
|
def test_runner_release_is_pinned(self):
|
|
self.assertEqual(controller.RUNNER_VERSION, "3.5.0")
|
|
self.assertEqual(len(controller.RUNNER_SHA256), 64)
|
|
|
|
def test_instance_spec_uses_pro6000_and_cuda_128(self):
|
|
spec = controller.instance_spec("image-test", 120, [])
|
|
self.assertEqual(spec["gpu_spec_uuid"], "pro6000-p")
|
|
self.assertEqual(spec["cuda_v_from"], 128)
|
|
self.assertNotIn("data_center_list", spec)
|
|
|
|
def test_instance_spec_preserves_data_center_preference(self):
|
|
spec = controller.instance_spec("image-test", 100, ["westDC3"])
|
|
self.assertEqual(spec["data_center_list"], ["westDC3"])
|
|
|
|
def test_instance_spec_rejects_missing_image(self):
|
|
with self.assertRaises(controller.ControllerError):
|
|
controller.instance_spec("", 120, [])
|
|
|
|
def test_public_image_catalog_has_unique_uuids(self):
|
|
uuids = [image["image_uuid"] for image in controller.PUBLIC_BASE_IMAGES]
|
|
self.assertEqual(len(uuids), len(set(uuids)))
|
|
|
|
def test_selects_exact_public_image_profile(self):
|
|
self.assertEqual(
|
|
controller.select_public_image("PyTorch", "2.0.0", "11.8"),
|
|
"base-image-l2t43iu6uk",
|
|
)
|
|
|
|
def test_refuses_to_downgrade_an_unknown_public_image_profile(self):
|
|
with self.assertRaises(controller.ControllerError):
|
|
controller.select_public_image("pytorch", "2.8.0", "12.8")
|
|
|
|
def test_runner_has_required_host_labels_and_no_docker(self):
|
|
rendered = controller.runner_yaml()
|
|
self.assertIn('"ubuntu-latest:host"', rendered)
|
|
self.assertIn('"ces-ci:host"', rendered)
|
|
self.assertIn('"gpu:host"', rendered)
|
|
self.assertIn('docker_host: "-"', rendered)
|
|
|
|
def test_state_is_private_and_round_trips(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
paths = controller.Paths(
|
|
root / "model",
|
|
"/model",
|
|
"/artifacts",
|
|
root / "artifacts",
|
|
root / "key",
|
|
root / "known_hosts",
|
|
root / "state.json",
|
|
root / "runner",
|
|
)
|
|
controller.write_state(paths, {"instance_uuid": "pro-test"})
|
|
self.assertEqual(controller.read_state(paths)["instance_uuid"], "pro-test")
|
|
self.assertEqual(os.stat(paths.state_file).st_mode & 0o777, 0o600)
|
|
|
|
def test_example_workflow_files_are_valid_json_free(self):
|
|
# Guards against accidentally treating workflow YAML as controller state.
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "value.json"
|
|
path.write_text(json.dumps({"ok": True}), encoding="utf-8")
|
|
self.assertTrue(json.loads(path.read_text())["ok"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|