Update tests to handle changes to _wmi_init

Crozzers · Nov 17, 2024 · 9f7c4c1 · 9f7c4c1
1 parent 680c47b
commit 9f7c4c1
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 25 deletions.
diff --git a/tests/mocks/windows_mock.py b/tests/mocks/windows_mock.py
@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 import re
 from collections import namedtuple
 from ctypes.wintypes import HMONITOR
@@ -110,8 +111,9 @@ def SetVCPFeature(handle, code, value_in):
             return 1
 
 
+@contextmanager
 def mock_wmi_init():
-    return FakeWMI()
+    yield FakeWMI()
 
 
 def mock_enum_display_devices():

diff --git a/tests/test_windows.py b/tests/test_windows.py
@@ -1,14 +1,21 @@
 import ctypes
-import screen_brightness_control as sbc
+from contextlib import contextmanager
 from typing import Type
-from unittest.mock import Mock
+from unittest.mock import Mock, call
+
 import pytest
 from pytest_mock import MockerFixture
-from unittest.mock import call
 
-from .helpers import BrightnessMethodTest
+import screen_brightness_control as sbc
 from screen_brightness_control.helpers import BrightnessMethod
-from .mocks.windows_mock import mock_enum_display_devices, mock_enum_display_monitors, mock_wmi_init, FakeWinDLL
+
+from .helpers import BrightnessMethodTest
+from .mocks.windows_mock import (
+    FakeWinDLL,
+    mock_enum_display_devices,
+    mock_enum_display_monitors,
+    mock_wmi_init,
+)
 
 
 @pytest.fixture
@@ -52,27 +59,35 @@ def test_without(self):
     class TestSetBrightness(BrightnessMethodTest.TestSetBrightness):
         class TestDisplayKwarg(BrightnessMethodTest.TestSetBrightness.TestDisplayKwarg):
             def test_with(self, mocker: MockerFixture, freeze_display_info, method, subtests):
-                wmi = sbc.windows._wmi_init()
-                mocker.patch.object(sbc.windows, '_wmi_init', Mock(return_value=wmi, spec=True))
-                brightness_method = wmi.WmiMonitorBrightnessMethods()[0]
-                mocker.patch.object(wmi, 'WmiMonitorBrightnessMethods', lambda: [brightness_method] * 3)
-                spy = mocker.spy(brightness_method, 'WmiSetBrightness')
-                for index, display in enumerate(freeze_display_info):
-                    with subtests.test(index=index):
-                        method.set_brightness(100, display=index)
-                        spy.assert_called_once_with(100, 0)
-                        spy.reset_mock()
+                with sbc.windows._wmi_init() as wmi:
+                    mocker.patch.object(
+                        sbc.windows,
+                        '_wmi_init',
+                        Mock(side_effect=contextmanager(lambda *_: (yield wmi)))
+                    )
+                    brightness_method = wmi.WmiMonitorBrightnessMethods()[0]
+                    mocker.patch.object(wmi, 'WmiMonitorBrightnessMethods', lambda: [brightness_method] * 3)
+                    spy = mocker.spy(brightness_method, 'WmiSetBrightness')
+                    for index, display in enumerate(freeze_display_info):
+                        with subtests.test(index=index):
+                            method.set_brightness(100, display=index)
+                            spy.assert_called_once_with(100, 0)
+                            spy.reset_mock()
 
             def test_without(self, mocker: MockerFixture, freeze_display_info, method):
-                wmi = sbc.windows._wmi_init()
-                mocker.patch.object(sbc.windows, '_wmi_init', Mock(return_value=wmi, spec=True))
-                brightness_method = wmi.WmiMonitorBrightnessMethods()[0]
-                mocker.patch.object(wmi, 'WmiMonitorBrightnessMethods', lambda: [brightness_method] * 3)
-                spy = mocker.spy(brightness_method, 'WmiSetBrightness')
-
-                method.set_brightness(100)
-                spy.assert_has_calls([call(100, 0)] * 3)
-                spy.reset_mock()
+                with sbc.windows._wmi_init() as wmi:
+                    mocker.patch.object(
+                        sbc.windows,
+                        '_wmi_init',
+                        Mock(side_effect=contextmanager(lambda *_: (yield wmi)))
+                    )
+                    brightness_method = wmi.WmiMonitorBrightnessMethods()[0]
+                    mocker.patch.object(wmi, 'WmiMonitorBrightnessMethods', lambda: [brightness_method] * 3)
+                    spy = mocker.spy(brightness_method, 'WmiSetBrightness')
+
+                    method.set_brightness(100)
+                    spy.assert_has_calls([call(100, 0)] * 3)
+                    spy.reset_mock()
 
 
 class TestVCP(BrightnessMethodTest):