This commit is contained in:
@@ -1,13 +1,40 @@
|
||||
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_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}
|
||||
|
||||
Reference in New Issue
Block a user