From 63db199bd6e624bd3f874d74bf88f9606d563420 Mon Sep 17 00:00:00 2001
From: katie <katie.may@canonical.com>
Date: Fri, 13 Dec 2024 16:30:42 +0100
Subject: [PATCH] tests: prepare memory limit before each test

---
 tests/lib/prepare-restore.sh |  7 +++++--
 tests/lib/prepare.sh         | 21 ++++++++++++---------
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/tests/lib/prepare-restore.sh b/tests/lib/prepare-restore.sh
index ed4a55f0484..e7eb6e55739 100755
--- a/tests/lib/prepare-restore.sh
+++ b/tests/lib/prepare-restore.sh
@@ -696,10 +696,13 @@ prepare_suite_each() {
         return 0
     fi
 
+    # shellcheck source=tests/lib/prepare.sh
+    . "$TESTSLIB"/prepare.sh
+    # Each individual task may potentially set the SNAP_NO_MEMORY_LIMIT variable
+    prepare_memory_limit_override
+
     if [[ "$variant" = full ]]; then
         if os.query is-classic; then
-            # shellcheck source=tests/lib/prepare.sh
-            . "$TESTSLIB"/prepare.sh
             prepare_each_classic
         fi
     fi
diff --git a/tests/lib/prepare.sh b/tests/lib/prepare.sh
index de17bf6461e..6e08cf4e040 100755
--- a/tests/lib/prepare.sh
+++ b/tests/lib/prepare.sh
@@ -248,14 +248,6 @@ update_core_snap_for_classic_reexec() {
 }
 
 prepare_memory_limit_override() {
-    # First time it is needed to save the initial env var value
-    if not tests.env is-set initial SNAPD_NO_MEMORY_LIMIT; then
-        tests.env set initial SNAPD_NO_MEMORY_LIMIT "$SNAPD_NO_MEMORY_LIMIT"
-    # Then if the new value is the same than the initial, then no new configuration needed
-    elif [ "$(tests.env get initial SNAPD_NO_MEMORY_LIMIT)" = "$SNAPD_NO_MEMORY_LIMIT" ]; then
-        return
-    fi
-
     # set up memory limits for snapd bu default unless explicit requested not to
     # or the system is known to be problematic
     local set_limit=1
@@ -279,6 +271,18 @@ prepare_memory_limit_override() {
             ;;
     esac
 
+    # If we don't wish to impose a memory limit, and the conf file 
+    # already doesn't exist, then no new configuration is needed
+    if [ "$set_limit" == "0" ] && ! [ -f "/etc/systemd/system/snapd.service.d/memory-max.conf" ]; then
+        return
+    fi
+
+    # If we wish to impose a memory limit, and the conf file 
+    # already exists, then no new configuration is needed
+    if [ "$set_limit" == "1" ] && [ -f "/etc/systemd/system/snapd.service.d/memory-max.conf" ]; then
+        return
+    fi
+
     if [ "$set_limit" = "0" ]; then
         # make sure the file does not exist then
         rm -f /etc/systemd/system/snapd.service.d/memory-max.conf
@@ -338,7 +342,6 @@ prepare_each_classic() {
     fi
 
     prepare_reexec_override
-    prepare_memory_limit_override
 }
 
 prepare_classic() {