Tests: Added --aggregate option to attach

Added a unit and an integration test for the --aggregate option. The unit test tests the case when --aggregate is given but --output is not. The --aggregate option is tested in the integration test.

Signed-off-by: Ivona Stojanovic <[email protected]>

tests/integration/test_attach.py

@@ -34,7 +34,7 @@ def bar():
 
 def baz():
     allocator = MemoryAllocator()
-    allocator.valloc(1024)
+    allocator.valloc(50 * 1024 * 1024)
     allocator.free()
 
 
@@ -52,13 +52,43 @@ def baz():
 """
 
 
-@pytest.mark.parametrize("method", ["lldb", "gdb"])
-def test_basic_attach(tmp_path, method):
-    if not debugger_available(method):
-        pytest.skip(f"a supported {method} debugger isn't installed")
+def compare_allocations(allocations1, allocations2):
+    assert len(allocations1) == len(allocations2)
+    for i in range(0, len(allocations1)):
+        assert allocations1[i].allocator == allocations2[i].allocator
+        assert allocations1[i].n_allocations == allocations2[i].n_allocations
+        assert allocations1[i].size == allocations2[i].size
+        assert allocations1[i].stack_id == allocations2[i].stack_id
+        assert allocations1[i].tid == allocations2[i].tid
+        assert allocations1[i].native_stack_id == allocations2[i].native_stack_id
+        assert (
+            allocations1[i].native_segment_generation
+            == allocations2[i].native_segment_generation
+        )
+        assert allocations1[i].thread_name == allocations2[i].thread_name
+
+
+def generate_command(method, output, aggregate):
+    cmd = [
+        sys.executable,
+        "-m",
+        "memray",
+        "attach",
+        "--verbose",
+        "--force",
+        "--method",
+        method,
+        "-o",
+        str(output),
+    ]
 
-    # GIVEN
-    output = tmp_path / "test.bin"
+    if aggregate:
+        cmd.append("--aggregate")
+
+    return cmd
+
+
+def run_process(cmd):
     tracked_process = subprocess.Popen(
         [sys.executable, "-uc", PROGRAM],
         stdin=subprocess.PIPE,
@@ -71,22 +101,12 @@ def test_basic_attach(tmp_path, method):
     assert tracked_process.stdout is not None
 
     assert tracked_process.stdout.readline() == "ready\n"
-    attach_cmd = [
-        sys.executable,
-        "-m",
-        "memray",
-        "attach",
-        "--verbose",
-        "--method",
-        method,
-        "-o",
-        str(output),
-        str(tracked_process.pid),
-    ]
+
+    cmd.append(str(tracked_process.pid))
 
     # WHEN
     try:
-        subprocess.check_output(attach_cmd, stderr=subprocess.STDOUT, text=True)
+        subprocess.check_output(cmd, stderr=subprocess.STDOUT, text=True)
     except subprocess.CalledProcessError as exc:
         if "Couldn't write extended state status" in exc.output:
             # https://bugs.debian.org/cgi-bin/bugreport.cgi?bug=898048
@@ -103,14 +123,65 @@ def test_basic_attach(tmp_path, method):
     assert "" == tracked_process.stdout.read()
     assert tracked_process.returncode == 0
 
-    reader = FileReader(output)
-    records = list(reader.get_allocation_records())
-    vallocs = [
+
+def get_functions(allocations):
+    (valloc,) = allocations
+    return [f[0] for f in valloc.stack_trace()]
+
+
+def get_relevant_allocations(records):
+    return [
         record
         for record in filter_relevant_allocations(records)
         if record.allocator == AllocatorType.VALLOC
     ]
 
-    (valloc,) = vallocs
-    functions = [f[0] for f in valloc.stack_trace()]
-    assert functions == ["valloc", "baz", "bar", "foo", "<module>"]
+
+@pytest.mark.parametrize("method", ["lldb", "gdb"])
+@pytest.mark.parametrize("aggregate", [True, False])
+def test_basic_attach(tmp_path, method, aggregate):
+    if not debugger_available(method):
+        pytest.skip(f"a supported {method} debugger isn't installed")
+
+    # GIVEN
+    output = tmp_path / "test.bin"
+
+    attach_cmd = generate_command(method, output, aggregate)
+
+    run_process(attach_cmd)
+
+    reader = FileReader(output)
+
+    # WHEN
+    try:
+        hwa_allocation_records = list(reader.get_high_watermark_allocation_records())
+        assert hwa_allocation_records is not None
+        allocation_records = list(reader.get_allocation_records())
+    except NotImplementedError as exc:
+        if aggregate:
+            assert (
+                "Can't get all allocations from a pre-aggregated capture file."
+                in str(exc)
+            )
+
+    hwa_relevant_allocations_records = get_relevant_allocations(hwa_allocation_records)
+    relevant_allocations_records = (
+        get_relevant_allocations(allocation_records) if not aggregate else []
+    )
+
+    if not aggregate:
+        assert get_functions(hwa_relevant_allocations_records) == get_functions(
+            relevant_allocations_records
+        )
+    else:
+        output_no_aggregate = tmp_path / "test.bin"
+        attach_cmd = generate_command(method, output_no_aggregate, False)
+        run_process(attach_cmd)
+
+        reader = FileReader(output_no_aggregate)
+        allocation_records = list(reader.get_high_watermark_allocation_records())
+        relevant_allocations_records = get_relevant_allocations(allocation_records)
+
+        compare_allocations(
+            relevant_allocations_records, hwa_relevant_allocations_records
+        )

tests/unit/test_attach.py

@@ -0,0 +1,21 @@
+from unittest.mock import patch
+
+import pytest
+
+from memray.commands import main
+
+
+@patch("memray.commands.attach.debugger_available")
+class TestAttachSubCommand:
+    def test_memray_attach_aggregated_without_output_file(
+        self, is_debugger_available_mock, capsys
+    ):
+        # GIVEN
+        is_debugger_available_mock.return_value = True
+
+        # WHEN
+        with pytest.raises(SystemExit):
+            main(["attach", "--aggregate", "1234"])
+
+        captured = capsys.readouterr()
+        assert "Can't use aggregated mode without an output file." in captured.err