forked from LeenkxTeam/Kmake
154 lines
4.5 KiB
Python
154 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright 2023 the V8 project authors. All rights reserved.
|
|
# Use of this source code is governed by a BSD-style license that can be
|
|
# found in the LICENSE file.
|
|
|
|
import os
|
|
import psutil
|
|
import threading
|
|
import unittest
|
|
|
|
from collections import namedtuple
|
|
from mock import patch
|
|
from pathlib import Path
|
|
from pyfakefs import fake_filesystem_unittest
|
|
|
|
import process_utils
|
|
|
|
|
|
Process = namedtuple('Process', 'pid')
|
|
MemoryInfo = namedtuple('MemoryInfo', 'rss, vms')
|
|
SystemMemory = namedtuple('SystemMemory', 'percent')
|
|
|
|
|
|
class TestStats(unittest.TestCase):
|
|
def test_update(self):
|
|
stats = process_utils.ProcessStats()
|
|
self.assertFalse(stats.available)
|
|
|
|
stats.update(MemoryInfo(17, 42))
|
|
self.assertTrue(stats.available)
|
|
self.assertEqual(17, stats.max_rss)
|
|
self.assertEqual(42, stats.max_vms)
|
|
|
|
stats.update(MemoryInfo(23, 23))
|
|
self.assertTrue(stats.available)
|
|
self.assertEqual(23, stats.max_rss)
|
|
self.assertEqual(42, stats.max_vms)
|
|
|
|
|
|
class TestProcessLogger(unittest.TestCase):
|
|
def test_null(self):
|
|
with process_utils.EmptyProcessLogger().log_stats(None) as stats:
|
|
pass
|
|
self.assertFalse(stats.available)
|
|
self.assertEqual(0, stats.max_rss)
|
|
self.assertEqual(0, stats.max_vms)
|
|
|
|
def test_basic(self):
|
|
"""Test three iterations of probing with mocked psutil."""
|
|
done = threading.Event()
|
|
results = [MemoryInfo(17, 101), MemoryInfo(2, 11), MemoryInfo(23, 42)]
|
|
|
|
class FakeProcess:
|
|
def __init__(self, process):
|
|
self.iter = iter(results)
|
|
|
|
def memory_info(self):
|
|
try:
|
|
return self.iter.__next__()
|
|
except StopIteration:
|
|
done.set()
|
|
raise psutil.NoSuchProcess(123)
|
|
|
|
logger = process_utils.PSUtilProcessLogger(0.01)
|
|
with patch('psutil.Process', FakeProcess):
|
|
with logger.log_stats(Process(123)) as stats:
|
|
done.wait()
|
|
|
|
self.assertTrue(stats.available)
|
|
self.assertEqual(23, stats.max_rss)
|
|
self.assertEqual(101, stats.max_vms)
|
|
|
|
def test_fast_process(self):
|
|
"""Test a process that finished fast."""
|
|
class FakeProcess:
|
|
def __init__(self, process):
|
|
raise psutil.NoSuchProcess(123)
|
|
|
|
logger = process_utils.PSUtilProcessLogger(0.01)
|
|
with patch('psutil.Process', FakeProcess):
|
|
with logger.log_stats(Process(123)) as stats:
|
|
pass
|
|
|
|
self.assertFalse(stats.available)
|
|
self.assertEqual(0, stats.max_rss)
|
|
self.assertEqual(0, stats.max_vms)
|
|
|
|
|
|
class TestLinuxProcessLogger(fake_filesystem_unittest.TestCase):
|
|
def test_child_process(self):
|
|
self.setUpPyfakefs(allow_root_user=True)
|
|
os.makedirs('/proc/123/task/123')
|
|
with open('/proc/123/task/123/children', 'w') as f:
|
|
f.write('42')
|
|
logger = process_utils.LinuxPSUtilProcessLogger()
|
|
self.assertEqual(42, logger.get_pid(123))
|
|
|
|
def test_no_child_process(self):
|
|
self.setUpPyfakefs(allow_root_user=True)
|
|
os.makedirs('/proc/123/task/123')
|
|
with open('/proc/123/task/123/children', 'w') as f:
|
|
f.write('')
|
|
logger = process_utils.LinuxPSUtilProcessLogger()
|
|
self.assertEqual(123, logger.get_pid(123))
|
|
|
|
def test_process_ended(self):
|
|
self.setUpPyfakefs(allow_root_user=True)
|
|
logger = process_utils.LinuxPSUtilProcessLogger()
|
|
self.assertEqual(123, logger.get_pid(123))
|
|
|
|
|
|
class TestSystemMemoryLogger(fake_filesystem_unittest.TestCase):
|
|
def test_log(self):
|
|
"""Test logging memory stats from a background thread.
|
|
|
|
Simulate deterministic stats and timestamps. The main thread waits
|
|
for at least two iterations.
|
|
"""
|
|
done = threading.Event()
|
|
results = 3 * [SystemMemory(1.7)]
|
|
it = iter(results)
|
|
|
|
def side_effect():
|
|
try:
|
|
return it.__next__()
|
|
except StopIteration:
|
|
done.set()
|
|
|
|
# If the main thread is slow, we might keep running into this.
|
|
return results[0]
|
|
|
|
output_file = Path('dir') / 'memory.log'
|
|
self.setUpPyfakefs()
|
|
self.fs.create_dir(output_file.parent)
|
|
logger = process_utils.PSUtilProcessLogger(0.01, 0.02)
|
|
|
|
with patch('psutil.virtual_memory', side_effect=side_effect):
|
|
with patch('time.time', return_value=13123212323):
|
|
with logger.log_system_memory(output_file):
|
|
done.wait()
|
|
|
|
# We get at least two lines of deterministic log. We might get more
|
|
# due to unpredictable delay after done.wait() on the main thread.
|
|
with open(output_file) as f:
|
|
lines = f.readlines()
|
|
expected = '2385-11-10 00:45:23 - 1.7%, 1.7%\n'
|
|
self.assertLess(1, len(lines))
|
|
self.assertEqual(expected, lines[0])
|
|
self.assertEqual(expected, lines[1])
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|