1.1 --- /dev/null Thu Jan 01 00:00:00 1970 +0000
1.2 +++ b/addon-sdk/source/python-lib/mozrunner/wpk.py Wed Dec 31 06:09:35 2014 +0100
1.3 @@ -0,0 +1,80 @@
1.4 +# This Source Code Form is subject to the terms of the Mozilla Public
1.5 +# License, v. 2.0. If a copy of the MPL was not distributed with this
1.6 +# file, You can obtain one at http://mozilla.org/MPL/2.0/.
1.7 +
1.8 +from ctypes import sizeof, windll, addressof, c_wchar, create_unicode_buffer
1.9 +from ctypes.wintypes import DWORD, HANDLE
1.10 +
1.11 +PROCESS_TERMINATE = 0x0001
1.12 +PROCESS_QUERY_INFORMATION = 0x0400
1.13 +PROCESS_VM_READ = 0x0010
1.14 +
1.15 +def get_pids(process_name):
1.16 + BIG_ARRAY = DWORD * 4096
1.17 + processes = BIG_ARRAY()
1.18 + needed = DWORD()
1.19 +
1.20 + pids = []
1.21 + result = windll.psapi.EnumProcesses(processes,
1.22 + sizeof(processes),
1.23 + addressof(needed))
1.24 + if not result:
1.25 + return pids
1.26 +
1.27 + num_results = needed.value / sizeof(DWORD)
1.28 +
1.29 + for i in range(num_results):
1.30 + pid = processes[i]
1.31 + process = windll.kernel32.OpenProcess(PROCESS_QUERY_INFORMATION |
1.32 + PROCESS_VM_READ,
1.33 + 0, pid)
1.34 + if process:
1.35 + module = HANDLE()
1.36 + result = windll.psapi.EnumProcessModules(process,
1.37 + addressof(module),
1.38 + sizeof(module),
1.39 + addressof(needed))
1.40 + if result:
1.41 + name = create_unicode_buffer(1024)
1.42 + result = windll.psapi.GetModuleBaseNameW(process, module,
1.43 + name, len(name))
1.44 + # TODO: This might not be the best way to
1.45 + # match a process name; maybe use a regexp instead.
1.46 + if name.value.startswith(process_name):
1.47 + pids.append(pid)
1.48 + windll.kernel32.CloseHandle(module)
1.49 + windll.kernel32.CloseHandle(process)
1.50 +
1.51 + return pids
1.52 +
1.53 +def kill_pid(pid):
1.54 + process = windll.kernel32.OpenProcess(PROCESS_TERMINATE, 0, pid)
1.55 + if process:
1.56 + windll.kernel32.TerminateProcess(process, 0)
1.57 + windll.kernel32.CloseHandle(process)
1.58 +
1.59 +if __name__ == '__main__':
1.60 + import subprocess
1.61 + import time
1.62 +
1.63 + # This test just opens a new notepad instance and kills it.
1.64 +
1.65 + name = 'notepad'
1.66 +
1.67 + old_pids = set(get_pids(name))
1.68 + subprocess.Popen([name])
1.69 + time.sleep(0.25)
1.70 + new_pids = set(get_pids(name)).difference(old_pids)
1.71 +
1.72 + if len(new_pids) != 1:
1.73 + raise Exception('%s was not opened or get_pids() is '
1.74 + 'malfunctioning' % name)
1.75 +
1.76 + kill_pid(tuple(new_pids)[0])
1.77 +
1.78 + newest_pids = set(get_pids(name)).difference(old_pids)
1.79 +
1.80 + if len(newest_pids) != 0:
1.81 + raise Exception('kill_pid() is malfunctioning')
1.82 +
1.83 + print "Test passed."
