71 lines
3.1 KiB
Python
71 lines
3.1 KiB
Python
|
import importlib.util
|
||
|
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
import tempfile
|
||
|
|
import unittest
|
||
|
|
from unittest.mock import patch
|
||
|
|
|
||
|
|
spec = importlib.util.spec_from_file_location(
|
||
|
|
'gpu_hook', Path(__file__).resolve().parents[1] / 'pve-shared-gpu-hook.py')
|
||
|
|
hook = importlib.util.module_from_spec(spec)
|
||
|
|
spec.loader.exec_module(hook)
|
||
|
|
|
||
|
|
|
||
|
|
class SharedGPUHookTests(unittest.TestCase):
|
||
|
|
def setUp(self):
|
||
|
|
self.directory = tempfile.TemporaryDirectory()
|
||
|
|
self.addCleanup(self.directory.cleanup)
|
||
|
|
root = Path(self.directory.name)
|
||
|
|
(root / 'boot').write_text('test-boot')
|
||
|
|
for name, value in [('STATE_DIR', root), ('BOOT_ID', root / 'boot')]:
|
||
|
|
p = patch.object(hook, name, value)
|
||
|
|
p.start()
|
||
|
|
self.addCleanup(p.stop)
|
||
|
|
|
||
|
|
def test_running_guest_or_busy_device_never_resets(self):
|
||
|
|
for failing_guard in ('stopped', 'unused_gpu'):
|
||
|
|
with self.subTest(guard=failing_guard), patch.object(hook, 'stopped'), \
|
||
|
|
patch.object(hook, 'unused_gpu'), patch.object(hook, 'reenumerate') as reset:
|
||
|
|
getattr(hook, failing_guard).side_effect = RuntimeError('busy')
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
hook.handle(102, 'pre-start')
|
||
|
|
reset.assert_not_called()
|
||
|
|
self.assertFalse((hook.STATE_DIR / 'reservation.json').exists())
|
||
|
|
|
||
|
|
def test_second_start_blocked_until_first_guest_stops(self):
|
||
|
|
with patch.object(hook, 'stopped'), patch.object(hook, 'unused_gpu'), \
|
||
|
|
patch.object(hook, 'reenumerate') as reset:
|
||
|
|
hook.handle(102, 'pre-start')
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
hook.handle(100, 'pre-start')
|
||
|
|
self.assertEqual(reset.call_count, 1)
|
||
|
|
hook.handle(100, 'post-stop')
|
||
|
|
self.assertTrue((hook.STATE_DIR / 'reservation.json').exists())
|
||
|
|
hook.handle(102, 'post-start')
|
||
|
|
hook.handle(102, 'post-stop')
|
||
|
|
hook.handle(100, 'pre-start')
|
||
|
|
self.assertEqual(reset.call_count, 2)
|
||
|
|
|
||
|
|
def test_failed_reset_stays_reserved_and_is_not_replayed(self):
|
||
|
|
with patch.object(hook, 'stopped'), patch.object(hook, 'unused_gpu'), \
|
||
|
|
patch.object(hook, 'reenumerate', side_effect=RuntimeError('device missing')) as reset:
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
hook.handle(102, 'pre-start')
|
||
|
|
hook.handle(102, 'post-stop')
|
||
|
|
state = json.loads((hook.STATE_DIR / 'reservation.json').read_text())
|
||
|
|
self.assertEqual(state['phase'], 'failed')
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
hook.handle(102, 'pre-start')
|
||
|
|
self.assertEqual(reset.call_count, 1)
|
||
|
|
|
||
|
|
def test_post_start_cannot_claim_another_guest_reservation(self):
|
||
|
|
with patch.object(hook, 'stopped'), patch.object(hook, 'unused_gpu'), \
|
||
|
|
patch.object(hook, 'reenumerate'):
|
||
|
|
hook.handle(102, 'pre-start')
|
||
|
|
with self.assertRaises(RuntimeError):
|
||
|
|
hook.handle(100, 'post-start')
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
unittest.main()
|