refactor(test/balloon): move logic to get guest avail mem to framework

Move the logic to get the MemAvailable from /proc/meminfo inside the
guest to a new guest_stats module in the test framework. This provides a
new class MeminfoGuest that can be used to retrieve this information
(and more!).

Signed-off-by: Riccardo Mancini <[email protected]>

Author: Manciukic

tests/framework/guest_stats.py

@@ -0,0 +1,79 @@
+# Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Classes for querying guest stats inside microVMs.
+"""
+
+
+class ByteUnit:
+    """Represents a byte unit that can be converted to other units."""
+
+    value_bytes: int
+
+    def __init__(self, value_bytes: int):
+        self.value_bytes = value_bytes
+
+    @classmethod
+    def from_kib(cls, value_kib: int):
+        """Creates a ByteUnit from a value in KiB."""
+        if value_kib < 0:
+            raise ValueError("value_kib must be non-negative")
+        return ByteUnit(value_kib * 1024)
+
+    def bytes(self) -> float:
+        """Returns the value in B."""
+        return self.value_bytes
+
+    def kib(self) -> float:
+        """Returns the value in KiB."""
+        return self.value_bytes / 1024
+
+    def mib(self) -> float:
+        """Returns the value in MiB."""
+        return self.value_bytes / (1 << 20)
+
+    def gib(self) -> float:
+        """Returns the value in GiB."""
+        return self.value_bytes / (1 << 30)
+
+
+class Meminfo:
+    """Represents the contents of /proc/meminfo inside the guest"""
+
+    mem_total: ByteUnit
+    mem_free: ByteUnit
+    mem_available: ByteUnit
+    buffers: ByteUnit
+    cached: ByteUnit
+
+    def __init__(self):
+        self.mem_total = ByteUnit(0)
+        self.mem_free = ByteUnit(0)
+        self.mem_available = ByteUnit(0)
+        self.buffers = ByteUnit(0)
+        self.cached = ByteUnit(0)
+
+
+class MeminfoGuest:
+    """Queries /proc/meminfo inside the guest"""
+
+    def __init__(self, vm):
+        self.vm = vm
+
+    def get(self) -> Meminfo:
+        """Returns the contents of /proc/meminfo inside the guest"""
+        meminfo = Meminfo()
+        for line in self.vm.ssh.check_output("cat /proc/meminfo").stdout.splitlines():
+            parts = line.split()
+            if parts[0] == "MemTotal:":
+                meminfo.mem_total = ByteUnit.from_kib(int(parts[1]))
+            elif parts[0] == "MemFree:":
+                meminfo.mem_free = ByteUnit.from_kib(int(parts[1]))
+            elif parts[0] == "MemAvailable:":
+                meminfo.mem_available = ByteUnit.from_kib(int(parts[1]))
+            elif parts[0] == "Buffers:":
+                meminfo.buffers = ByteUnit.from_kib(int(parts[1]))
+            elif parts[0] == "Cached:":
+                meminfo.cached = ByteUnit.from_kib(int(parts[1]))
+
+        return meminfo

tests/framework/utils.py

@@ -254,25 +254,6 @@ def search_output_from_cmd(cmd: str, find_regex: typing.Pattern) -> typing.Match
     )
 
 
-def get_free_mem_ssh(ssh_connection):
-    """
-    Get how much free memory in kB a guest sees, over ssh.
-
-    :param ssh_connection: connection to the guest
-    :return: available mem column output of 'free'
-    """
-    _, stdout, stderr = ssh_connection.run("cat /proc/meminfo | grep MemAvailable")
-    assert stderr == ""
-
-    # Split "MemAvailable:   123456 kB" and validate it
-    meminfo_data = stdout.split()
-    if len(meminfo_data) == 3:
-        # Return the middle element in the array
-        return int(meminfo_data[1])
-
-    raise Exception("Available memory not found in `/proc/meminfo")
-
-
 def _format_output_message(proc, stdout, stderr):
     output_message = f"\n[{proc.pid}] Command:\n{proc.args}"
     # Append stdout/stderr to the output message

tests/integration_tests/functional/test_balloon.py

@@ -9,6 +9,7 @@
 import pytest
 import requests
 
+from framework.guest_stats import MeminfoGuest
 from framework.utils import get_resident_memory
 
 STATS_POLLING_INTERVAL_S = 1
@@ -142,18 +143,18 @@ def test_inflate_reduces_free(uvm_plain_any):
 
     # Start the microvm
     test_microvm.start()
-    firecracker_pid = test_microvm.firecracker_pid
+    meminfo = MeminfoGuest(test_microvm)
 
     # Get the free memory before ballooning.
-    available_mem_deflated = get_free_mem_ssh(test_microvm.ssh)
+    available_mem_deflated = meminfo.get().mem_free.kib()
 
     # Inflate 64 MB == 16384 page balloon.
     test_microvm.api.balloon.patch(amount_mib=64)
     # This call will internally wait for rss to become stable.
-    _ = get_stable_rss_mem(test_microvm.ps)
+    _ = get_stable_rss_mem(test_microvm)
 
     # Get the free memory after ballooning.
-    available_mem_inflated = get_free_mem_ssh(test_microvm.ssh)
+    available_mem_inflated = meminfo.get().mem_free.kib()
 
     # Assert that ballooning reclaimed about 64 MB of memory.
     assert available_mem_inflated <= available_mem_deflated - 85 * 64000 / 100