blender-benchmark-bundle/tests/test_foundation_system_info.py

import json
import sys
import unittest
from unittest import mock
import typing

from benchmark.foundation import system_info


class FakeCycles:
    """Injected as _cycles module so we can mock its results."""

    @staticmethod
    def available_devices():
        return None


class FakeDevice:
    def __init__(self, name: str, type: str, *, use: bool):
        self.name = name
        self.type = type
        self.use = use

    def as_cycles(self) -> typing.Tuple[str, str]:
        return self.name, self.type


class FakeCpref:
    def __init__(self, compute_device_type: str, devices: typing.List[FakeDevice]):
        assert compute_device_type in {'NONE', 'CPU', 'CUDA', 'OPENCL'}
        self.compute_device_type = compute_device_type
        self.devices = devices


class AbstractFakeDevicesTest(unittest.TestCase):
    cpu = 'Intel Core i7-4790K CPU @ 4.00GHz'
    gpu = 'GeForce GTX 970'  # display
    gpu2 = 'GeForce GTX 980'  # non-display

    compute_devices_as_printed = [
        {'name': cpu, 'type': 'CPU'},
        {'name': gpu, 'type': 'CUDA', 'is_display': True},
        {'name': gpu2, 'type': 'CUDA', 'is_display': False},
    ]

    cpref_devs = [
        FakeDevice(cpu, 'CPU', use=True),
        FakeDevice(f'{gpu} (Display)', 'CUDA', use=True),
        FakeDevice(gpu2, 'CUDA', use=False),
    ]

    def setUp(self):
        super().setUp()

        sys.modules['bpy'] = sys.modules[__name__]  # don't fail on 'import bpy'
        sys.modules['_cycles'] = FakeCycles  # don't fail on 'import _cycles'

        # Import so that we can mock functions here.
        # noinspection PyUnresolvedReferences
        from benchmark import configure

    def expected_benchmark_devices_output(self) -> str:
        return json.dumps(self.compute_devices_as_printed, sort_keys=True)


class ConfigurePyTest(AbstractFakeDevicesTest):

    @mock.patch('_cycles.available_devices')
    @mock.patch('benchmark.configure.print')
    def test_log_compute_devices_cpu(self, mock_print, mock_available_devices):
        from benchmark import configure

        mock_available_devices.return_value = [dev.as_cycles() for dev in self.cpref_devs]
        cpref = FakeCpref(compute_device_type='NONE', devices=self.cpref_devs)
        configure.logComputeDevices(cpref)

        expect_json = json.dumps({'name': self.cpu}, sort_keys=True)
        mock_print.assert_has_calls([
            mock.call('Compute device type: CPU'),
            mock.call(f'Using compute device: {expect_json}'),
        ])

    @mock.patch('_cycles.available_devices')
    @mock.patch('benchmark.configure.print')
    def test_log_compute_devices_gpu(self, mock_print, mock_available_devices):
        from benchmark import configure

        mock_available_devices.return_value = [dev.as_cycles() for dev in self.cpref_devs]

        cpref = FakeCpref(compute_device_type='CUDA', devices=self.cpref_devs)
        configure.logComputeDevices(cpref)

        expect_json = json.dumps({'name': self.gpu, 'is_display': True}, sort_keys=True)
        mock_print.assert_has_calls([
            mock.call('Compute device type: CUDA'),
            mock.call(f'Using compute device: {expect_json}'),
        ])

    @mock.patch('benchmark.configure.print')
    def test_log_system_info(self, mock_print):
        from benchmark import configure

        cpref = FakeCpref(compute_device_type='CUDA', devices=self.cpref_devs)
        configure.logSystemInfo(cpref)

        mock_print.assert_has_calls([
            mock.call(f'Benchmark Devices: {self.expected_benchmark_devices_output()}'),
        ])


class BenchRunnerTest(AbstractFakeDevicesTest):

    @mock.patch('subprocess.Popen')
    def test_get_device_info_gpu(self, mock_popen):
        send_json = json.dumps({'name': self.gpu, 'is_display': True}, sort_keys=True)
        blender_output = ('\n'.join([
            'Nonsense lines',
            'Compute device type: GPU',
            f'Using compute device: {send_json}',
            'CPU threads used: 47',
            'More nonsense lines',
        ])).encode()

        mock_process = mock.Mock()
        mock_popen.return_value = mock_process
        mock_process.communicate.return_value = blender_output, b''

        from benchmark.foundation import benchrunner, context

        ctx = context.Context()
        info = benchrunner.benchmarkGetDeviceInfo(ctx)

        expected_info = {
            "device_type": 'GPU',
            "compute_devices": [{'name': self.gpu, 'is_display': True}],
            "num_cpu_threads": 47
        }
        self.assertEqual(expected_info, info)