From d6064c4d3d6d33885c27a3642566316ecf0583f3 Mon Sep 17 00:00:00 2001
From: Xichen Lin <lukelin0907@gmail.com>
Date: Thu, 18 Aug 2022 04:48:09 +0000
Subject: [PATCH] update unit tests for swap allocator

Signed-off-by: Xichen Lin <lukelin0907@gmail.com>
---
 tests/swap_allocator_test.py | 64 ++++++++++++++++++++++++++++++++++++
 1 file changed, 64 insertions(+)

diff --git a/tests/swap_allocator_test.py b/tests/swap_allocator_test.py
index 033b215dd23..960d4e8caf9 100644
--- a/tests/swap_allocator_test.py
+++ b/tests/swap_allocator_test.py
@@ -17,6 +17,8 @@ def test_read_from_meminfo(self):
         proc_meminfo_lines = [
             "MemTotal:       32859496 kB",
             "MemFree:        16275512 kB",
+            "SwapTotal:      2000000  kB",
+            "SwapFree:       1000000  kB",
             "HugePages_Total:       0",
             "HugePages_Free:        0",
         ]
@@ -24,6 +26,8 @@ def test_read_from_meminfo(self):
         read_meminfo_expected_return = {
             "MemTotal": 32859496,
             "MemFree": 16275512,
+            "SwapTotal": 2000000,
+            "SwapFree": 1000000,
             "HugePages_Total": 0,
             "HugePages_Free": 0
         }
@@ -113,6 +117,8 @@ def test_swap_allocator_context_enter_allocate_true_insufficient_total_memory(se
             mock_meminfo.return_value = {
                 "MemTotal": 2000000,
                 "MemAvailable": 1900000,
+                "SwapTotal": 0,
+                "SwapFree": 0,
             }
             mock_exists.return_value = False
 
@@ -135,6 +141,56 @@ def test_swap_allocator_context_enter_allocate_true_insufficient_available_memor
             mock_meminfo.return_value = {
                 "MemTotal": 3000000,
                 "MemAvailable": 1000000,
+                "SwapTotal": 0,
+                "SwapFree": 0,
+            }
+            mock_exists.return_value = False
+
+            swap_allocator = SWAPAllocator(allocate=True)
+            try:
+                swap_allocator.__enter__()
+            except Exception as detail:
+                pytest.fail("SWAPAllocator context manager should not raise exception %s" % repr(detail))
+            mock_setup.assert_called_once()
+            mock_remove.assert_not_called()
+            assert swap_allocator.is_allocated is True
+
+    def test_swap_allocator_context_enter_allocate_true_insufficient_total_memory_plus_swap(self):
+        with mock.patch("sonic_installer.main.SWAPAllocator.get_disk_freespace") as mock_disk_free, \
+                mock.patch("sonic_installer.main.SWAPAllocator.read_from_meminfo") as mock_meminfo, \
+                mock.patch("sonic_installer.main.SWAPAllocator.setup_swapmem") as mock_setup, \
+                mock.patch("sonic_installer.main.SWAPAllocator.remove_swapmem") as mock_remove, \
+                mock.patch("os.path.exists") as mock_exists:
+            mock_disk_free.return_value = 10 * 1024 * 1024 * 1024
+            mock_meminfo.return_value = {
+                "MemTotal": 1000000,
+                "MemAvailable": 900000,
+                "SwapTotal": 1000000,
+                "SwapFree": 1000000,
+            }
+            mock_exists.return_value = False
+
+            swap_allocator = SWAPAllocator(allocate=True)
+            try:
+                swap_allocator.__enter__()
+            except Exception as detail:
+                pytest.fail("SWAPAllocator context manager should not raise exception %s" % repr(detail))
+            mock_setup.assert_called_once()
+            mock_remove.assert_not_called()
+            assert swap_allocator.is_allocated is True
+
+    def test_swap_allocator_context_enter_allocate_true_insufficient_available_memory_plus_swap(self):
+        with mock.patch("sonic_installer.main.SWAPAllocator.get_disk_freespace") as mock_disk_free, \
+                mock.patch("sonic_installer.main.SWAPAllocator.read_from_meminfo") as mock_meminfo, \
+                mock.patch("sonic_installer.main.SWAPAllocator.setup_swapmem") as mock_setup, \
+                mock.patch("sonic_installer.main.SWAPAllocator.remove_swapmem") as mock_remove, \
+                mock.patch("os.path.exists") as mock_exists:
+            mock_disk_free.return_value = 10 * 1024 * 1024 * 1024
+            mock_meminfo.return_value = {
+                "MemTotal": 2000000,
+                "MemAvailable": 500000,
+                "SwapTotal": 1000000,
+                "SwapFree": 500000,
             }
             mock_exists.return_value = False
 
@@ -157,6 +213,8 @@ def test_swap_allocator_context_enter_allocate_true_insufficient_disk_space(self
             mock_meminfo.return_value = {
                 "MemTotal": 32859496,
                 "MemAvailable": 16275512,
+                "SwapTotal": 0,
+                "SwapFree": 0,
             }
             mock_exists.return_value = False
 
@@ -179,6 +237,8 @@ def test_swap_allocator_context_enter_allocate_true_swapfile_present(self):
             mock_meminfo.return_value = {
                 "MemTotal": 32859496,
                 "MemAvailable": 1000000,
+                "SwapTotal": 0,
+                "SwapFree": 0,
             }
             mock_exists.return_value = True
 
@@ -201,6 +261,8 @@ def test_swap_allocator_context_enter_setup_error(self):
             mock_meminfo.return_value = {
                 "MemTotal": 32859496,
                 "MemAvailable": 1000000,
+                "SwapTotal": 0,
+                "SwapFree": 0,
             }
             mock_exists.return_value = False
             expected_err_str = "Pseudo Error"
@@ -225,6 +287,8 @@ def test_swap_allocator_context_enter_allocate_false(self):
             mock_meminfo.return_value = {
                 "MemTotal": 32859496,
                 "MemAvailable": 1000000,
+                "SwapTotal": 0,
+                "SwapFree": 0,
             }
             mock_exists.return_value = False