| # Copyright 2026 The Chromium Authors |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| """ Tests for memory_usage_view.py""" |
| |
| import json |
| import os |
| import pandas as pd |
| import pathlib |
| import sys |
| import unittest |
| |
| _SRC_PATH = pathlib.Path(__file__).resolve().parents[4] |
| sys.path.append(str(_SRC_PATH / 'tools/android')) |
| from colabutils.memory_usage.memory_usage_view import ( |
| MemoryUsageView, TreeNode, _aggregate_nodes, _zip_by_name, |
| _compare_node_lists, _prettify_size) |
| |
| _REALISTIC_JSON = """ |
| [ |
| { |
| "name": "main", |
| "value": 0, |
| "children": [ |
| { |
| "name": "foo", |
| "value": 0, |
| "children": [ |
| { |
| "name": "malloc", |
| "value": 8 |
| } |
| ] |
| }, |
| { |
| "name": "malloc", |
| "value": 16 |
| }, |
| { |
| "name": "bar", |
| "value": 0, |
| "children": [ |
| { |
| "name": "malloc", |
| "value": 48 |
| } |
| ] |
| } |
| ] |
| } |
| ] |
| """ |
| |
| _SIMPLE_JSON_WITH_DELTAS = """ |
| [ |
| { |
| "name": "root1", |
| "value": 100, |
| "delta": 10, |
| "children": [ |
| { |
| "name": "child1_1", |
| "value": 50, |
| "delta": 5 |
| }, |
| { |
| "name": "child1_2", |
| "value": 30, |
| "delta": 3 |
| } |
| ] |
| }, |
| { |
| "name": "root2", |
| "value": 200, |
| "delta": 20 |
| } |
| ] |
| """ |
| |
| |
| class MemoryUsageViewTest(unittest.TestCase): |
| |
| def test_json_preserved(self): |
| view = MemoryUsageView.from_json(_SIMPLE_JSON_WITH_DELTAS) |
| dumped_json = view.to_json() |
| self.assertEqual(json.loads(_SIMPLE_JSON_WITH_DELTAS), |
| json.loads(dumped_json)) |
| |
| def test_from_json_defaults(self): |
| json_data = '''[{"name": "root"}]''' |
| view = MemoryUsageView.from_json(json_data) |
| self.assertEqual(len(view.roots), 1) |
| self.assertEqual(view.roots[0].name, 'root') |
| self.assertEqual(view.roots[0].value, 0) |
| self.assertEqual(view.roots[0].delta, 0) |
| self.assertEqual(view.roots[0].children, []) |
| |
| def test_from_json_empty_name_stays_empty(self): |
| json_data = '''[{"name": ""}]''' |
| view = MemoryUsageView.from_json(json_data) |
| self.assertEqual(len(view.roots), 1) |
| self.assertEqual(view.roots[0].name, '') |
| |
| def test_to_json_no_children(self): |
| node = TreeNode(name='test_node') |
| view = MemoryUsageView(roots=[node]) |
| dumped_json = json.loads(view.to_json()) |
| expected_json = {'name': 'test_node', 'value': 0, 'delta': 0} |
| self.assertEqual(dumped_json, [expected_json]) |
| |
| def test_from_df(self): |
| df_data = { |
| 'callsite_id': [1, 2, 3], |
| 'depth': [0, 0, 1], |
| 'frame_name': ['root_a', 'root_b', 'child_a_1'], |
| 'parent_callsite_id': [None, None, 1], |
| 'total_size_bytes': [100, 200, 50], |
| } |
| df = pd.DataFrame(df_data) |
| view = MemoryUsageView.from_df(df, demangle=False, aggregate=False) |
| dumped_json = view.to_json() |
| expected_json = """ |
| [ |
| { |
| "name": "root_a", |
| "value": 100, |
| "delta": 0, |
| "children": [ |
| { |
| "name": "child_a_1", |
| "value": 50, |
| "delta": 0 |
| } |
| ] |
| }, |
| { |
| "name": "root_b", |
| "value": 200, |
| "delta": 0 |
| } |
| ] |
| """ |
| self.assertEqual(json.loads(expected_json), json.loads(dumped_json)) |
| self.assertEqual(['root_a', 'root_b'], sorted(view.toplevel_names())) |
| |
| def test_from_df_empty_name_converts_to_unknown(self): |
| df_data = { |
| 'callsite_id': [1, 2, 3], |
| 'depth': [0, 0, 1], |
| 'frame_name': ['root_a', 'root_b', ''], |
| 'parent_callsite_id': [None, None, 1], |
| 'total_size_bytes': [100, 200, 50], |
| } |
| df = pd.DataFrame(df_data) |
| view = MemoryUsageView.from_df(df, demangle=False, aggregate=True) |
| self.assertEqual(2, len(view.roots)) |
| root_b_node, root_a_node = view.roots |
| self.assertEqual('root_a', root_a_node.name) |
| self.assertEqual(100 + 50, root_a_node.value) |
| self.assertEqual(1, len(root_a_node.children)) |
| empty_name_node = root_a_node.children[0] |
| self.assertEqual('unknown', empty_name_node.name) |
| self.assertEqual(50, empty_name_node.value) |
| |
| def test_aggregation_sorting(self): |
| # Check all cumulative values and the order. |
| view = MemoryUsageView.from_json(_REALISTIC_JSON) |
| view.roots = _aggregate_nodes(view.roots) |
| self.assertEqual(1, len(view.roots)) |
| root_node = view.roots[0] |
| self.assertEqual('main', root_node.name) |
| self.assertEqual(48 + 16 + 8, root_node.value) |
| self.assertEqual(3, len(root_node.children)) |
| bar_node, malloc_node, foo_node = root_node.children |
| self.assertEqual('bar', bar_node.name) |
| self.assertEqual(48, bar_node.value) |
| self.assertEqual('malloc', malloc_node.name) |
| self.assertEqual(16, malloc_node.value) |
| self.assertEqual('foo', foo_node.name) |
| self.assertEqual(8, foo_node.value) |
| |
| def test_aggregation_sorting2(self): |
| # Change the order of the children and verify that |
| # it leads to the same order after aggregation. |
| view = MemoryUsageView.from_json(_REALISTIC_JSON) |
| foo_node, malloc_node, bar_node = view.roots[0].children |
| view.roots[0].children = [foo_node, bar_node, malloc_node] |
| view.roots = _aggregate_nodes(view.roots) |
| root_node = view.roots[0] |
| bar_node, malloc_node, foo_node = root_node.children |
| self.assertEqual('bar', bar_node.name) |
| self.assertEqual(48, bar_node.value) |
| self.assertEqual('malloc', malloc_node.name) |
| self.assertEqual(16, malloc_node.value) |
| self.assertEqual('foo', foo_node.name) |
| self.assertEqual(8, foo_node.value) |
| |
| def test_aggregation_merge_duplicate_names(self): |
| view = MemoryUsageView.from_json(_REALISTIC_JSON) |
| # Make two adjacent node names match for merging. |
| foo_node, _, bar_node = view.roots[0].children |
| self.assertEqual('bar', bar_node.name) |
| foo_node.name = 'bar' |
| view.roots = _aggregate_nodes(view.roots) |
| self.assertEqual(1, len(view.roots)) |
| root_node = view.roots[0] |
| self.assertEqual(2, len(root_node.children)) |
| bar_node, malloc_node = root_node.children |
| self.assertEqual('bar', bar_node.name) |
| self.assertEqual(48 + 8, bar_node.value) |
| self.assertEqual('malloc', malloc_node.name) |
| self.assertEqual(16, malloc_node.value) |
| |
| def test_zip_by_name_both_empty(self): |
| self.assertEqual([], list(_zip_by_name([], []))) |
| |
| def test_zip_by_name_one_empty(self): |
| node = TreeNode(name='abc') |
| self.assertEqual([(node, None)], list(_zip_by_name([node], []))) |
| self.assertEqual([(None, node)], list(_zip_by_name([], [node]))) |
| node2 = TreeNode(name='def') |
| self.assertEqual([(None, node), (None, node2)], |
| list(_zip_by_name([], [node, node2]))) |
| |
| def test_zip_by_name_simple(self): |
| left_node1 = TreeNode(name='l1') |
| left_node2 = TreeNode(name='l2') |
| left_nodes = [left_node1, left_node2] |
| right_node1 = TreeNode(name='r1') |
| right_node2 = TreeNode(name='r2') |
| right_nodes = [right_node1, right_node2] |
| |
| result = list(_zip_by_name(left_nodes, right_nodes)) |
| |
| self.assertEqual(4, len(result)) |
| self.assertIn((left_node1, None), result) |
| self.assertIn((left_node2, None), result) |
| self.assertIn((None, right_node1), result) |
| self.assertIn((None, right_node2), result) |
| |
| def test_zip_by_name_with_intersection(self): |
| left_node1 = TreeNode(name='l1') |
| left_abc = TreeNode(name='abc') |
| left_nodes = [left_node1, left_abc] |
| right_abc = TreeNode(name='abc') |
| right_node2 = TreeNode(name='r2') |
| right_nodes = [right_abc, right_node2] |
| |
| result = list(_zip_by_name(left_nodes, right_nodes)) |
| |
| self.assertEqual(3, len(result)) |
| self.assertIn((left_abc, right_abc), result) |
| self.assertIn((left_node1, None), result) |
| self.assertIn((None, right_node2), result) |
| |
| def test_compare_nothing(self): |
| self.assertEqual([], _compare_node_lists([], [])) |
| |
| def test_compare_one_node_to_nothing(self): |
| node = TreeNode(name='one', value=123) |
| result_node = TreeNode(name='one', value=123, delta=-123) |
| |
| self.assertEqual([result_node], _compare_node_lists([node], [])) |
| |
| node.value = 456 |
| result_node.value = 0 |
| result_node.delta = 456 |
| |
| self.assertEqual([result_node], _compare_node_lists([], [node])) |
| |
| def test_compare_same_name_nodes(self): |
| node_base = TreeNode(name='pikachu', value=123) |
| node_new = TreeNode(name='pikachu', value=456) |
| |
| result = _compare_node_lists([node_base], [node_new]) |
| self.assertEqual(1, len(result)) |
| result_node = result[0] |
| self.assertEqual(123, result_node.value) |
| self.assertEqual('pikachu', result_node.name) |
| self.assertEqual(456 - 123, result_node.delta) |
| |
| def test_compare_children(self): |
| base_root = TreeNode(name='root', value=30) |
| new_root = TreeNode(name='root', value=80) |
| base_root.children = [ |
| TreeNode(name='common_child', value=10), |
| TreeNode(name='base_only_child', value=20), |
| ] |
| new_root.children = [ |
| TreeNode(name='common_child', value=30), |
| TreeNode(name='new_only_child', value=50), |
| ] |
| |
| result_nodes = _compare_node_lists([base_root], [new_root]) |
| self.assertEqual(1, len(result_nodes)) |
| result_root = result_nodes[0] |
| |
| self.assertEqual('root', result_root.name) |
| self.assertEqual(30, result_root.value) |
| self.assertEqual(80 - 30, result_root.delta) |
| |
| self.assertEqual(3, len(result_root.children)) |
| new_only, common_child, base_only = result_root.children |
| |
| self.assertEqual('new_only_child', new_only.name) |
| self.assertEqual(0, new_only.value) |
| self.assertEqual(50, new_only.delta) |
| |
| self.assertEqual('common_child', common_child.name) |
| self.assertEqual(10, common_child.value) |
| self.assertEqual(30 - 10, common_child.delta) |
| |
| self.assertEqual('base_only_child', base_only.name) |
| self.assertEqual(20, base_only.value) |
| self.assertEqual(-20, base_only.delta) |
| |
| def test_compare_children_and_single(self): |
| base_root = TreeNode(name='root', value=30) |
| new_root = TreeNode(name='root', value=80) |
| base_root.children = [ |
| TreeNode(name='child10', value=10), |
| TreeNode(name='child20', value=20), |
| ] |
| result_nodes = _compare_node_lists([base_root], [new_root]) |
| self.assertEqual(1, len(result_nodes)) |
| result_root = result_nodes[0] |
| |
| self.assertEqual('root', result_root.name) |
| self.assertEqual(30, result_root.value) |
| self.assertEqual(80 - 30, result_root.delta) |
| |
| self.assertEqual(2, len(result_root.children)) |
| child_10, child_20 = result_root.children |
| self.assertEqual(10, child_10.value) |
| self.assertEqual(-10, child_10.delta) |
| self.assertEqual(20, child_20.value) |
| self.assertEqual(-20, child_20.delta) |
| |
| def test_prettify_size(self): |
| tests = [ |
| (500, 500), |
| (1024, '1 KiB'), |
| (1536, '1.5 KiB'), |
| (1500, '1.46 KiB'), |
| (1048576, '1 MiB'), |
| (123456789, '117.74 MiB'), |
| (1099511627776, '1024 GiB'), |
| (2000000000000, 2000000000000), |
| ] |
| |
| test_lines = [] |
| for inp, expected in tests: |
| self.assertEqual(expected, _prettify_size(inp)) |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
