test the internal wmi module on windows this is used by the platform module and potentially others do this first so test will be skipped if module doesn t exist best we can check for the version is that it s digits dot digits anything otherwise we are likely checking the result of the query against itself repeated queries should not break invalid queries fail with oserror this is the expected error code all others should fail the test queries other than select are blocked to avoid potential exploits ensure very big queries fail test multiple times to ensure consistency multiple instances should have an extra null separator test the internal _wmi module on windows this is used by the platform module and potentially others do this first so test will be skipped if module doesn t exist best we can check for the version is that it s digits dot digits anything otherwise we are likely checking the result of the query against itself repeated queries should not break invalid queries fail with oserror this is the expected error code all others should fail the test queries other than select are blocked to avoid potential exploits ensure very big queries fail test multiple times to ensure consistency multiple instances should have an extra null separator
import unittest
from test.support import import_helper, requires_resource



_wmi = import_helper.import_module('_wmi', required_on=['win'])


class WmiTests(unittest.TestCase):
    def test_wmi_query_os_version(self):
        r = _wmi.exec_query("SELECT Version FROM Win32_OperatingSystem").split("\0")
        self.assertEqual(1, len(r))
        k, eq, v = r[0].partition("=")
        self.assertEqual("=", eq, r[0])
        self.assertEqual("Version", k, r[0])
        
        
        self.assertRegex(v, r"\d+\.\d+.+$", r[0])

    def test_wmi_query_repeated(self):
        
        for _ in range(10):
            self.test_wmi_query_os_version()

    def test_wmi_query_error(self):
        
        try:
            _wmi.exec_query("SELECT InvalidColumnName FROM InvalidTableName")
        except OSError as ex:
            if ex.winerror & 0xFFFFFFFF == 0x80041010:
                
                return
        self.fail("Expected OSError")

    def test_wmi_query_repeated_error(self):
        for _ in range(10):
            self.test_wmi_query_error()

    def test_wmi_query_not_select(self):
        
        with self.assertRaises(ValueError):
            _wmi.exec_query("not select, just in case someone tries something")

    @requires_resource('cpu')
    def test_wmi_query_overflow(self):
        
        
        for _ in range(2):
            with self.assertRaises(OSError):
                _wmi.exec_query("SELECT * FROM CIM_DataFile")

    def test_wmi_query_multiple_rows(self):
        
        r = _wmi.exec_query("SELECT ProcessId FROM Win32_Process WHERE ProcessId < 1000")
        self.assertFalse(r.startswith("\0"), r)
        self.assertFalse(r.endswith("\0"), r)
        it = iter(r.split("\0"))
        try:
            while True:
                self.assertRegex(next(it), r"ProcessId=\d+")
                self.assertEqual("", next(it))
        except StopIteration:
            pass

    def test_wmi_query_threads(self):
        from concurrent.futures import ThreadPoolExecutor
        query = "SELECT ProcessId FROM Win32_Process WHERE ProcessId < 1000"
        with ThreadPoolExecutor(4) as pool:
            task = [pool.submit(_wmi.exec_query, query) for _ in range(32)]
            for t in task:
                self.assertRegex(t.result(), "ProcessId=")