From b9a6049a954f6053b49de198bbacf550d5728de7 Mon Sep 17 00:00:00 2001
From: Changrong Wu <wcr@live.cn>
Date: Tue, 16 Jul 2024 15:53:16 -0700
Subject: [PATCH] [Bug Fix] Fix disk check test and drops group test (#3424)

* tests/disk_check_test.py: remove temp files during teardown

- modify teardown_class() to remove /tmp/tmp*

* tests/drops_group_test.py: add code to remove temporary files when setting up test class

- add a remove_tmp_dropstat_file() function as a helper to clean the cache
- add an invocation of remove_tmp_dropstat_file() in setup_class() of TestDropCounters class

* tests/disk_check_test.py: fix the subprocess command in the teardown_class() function

* tests/disk_check_test.py: fix formatting for pre-commit check

* tests/drops_group_test.py: fix formatting for pre-commit check
---
 tests/disk_check_test.py  |  5 ++++-
 tests/drops_group_test.py | 10 +++++++---
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/tests/disk_check_test.py b/tests/disk_check_test.py
index 82b8b16ff6..ac541b05b9 100644
--- a/tests/disk_check_test.py
+++ b/tests/disk_check_test.py
@@ -1,7 +1,6 @@
 import sys
 import syslog
 from unittest.mock import patch
-import pytest
 import subprocess
 
 sys.path.append("scripts")
@@ -178,3 +177,7 @@ def test_readonly(self, mock_proc, mock_log):
             
         assert max_log_lvl == syslog.LOG_ERR
 
+    @classmethod
+    def teardown_class(cls):
+        subprocess.run("rm -rf /tmp/tmp*", shell=True)  # cleanup the temporary dirs
+        print("TEARDOWN")
diff --git a/tests/drops_group_test.py b/tests/drops_group_test.py
index ad8c8a4203..93f99e3f1b 100644
--- a/tests/drops_group_test.py
+++ b/tests/drops_group_test.py
@@ -3,6 +3,7 @@
 
 import shutil
 from click.testing import CliRunner
+from utilities_common.cli import UserCache
 
 test_path = os.path.dirname(os.path.abspath(__file__))
 modules_path = os.path.dirname(test_path)
@@ -97,14 +98,17 @@
 sonic_drops_test               0                    0
 """
 
-dropstat_path = "/tmp/dropstat-27"
+
+def remove_tmp_dropstat_file():
+    # remove the tmp portstat
+    cache = UserCache("dropstat")
+    cache.remove_all()
 
 class TestDropCounters(object):
     @classmethod
     def setup_class(cls):
         print("SETUP")
-        if os.path.exists(dropstat_path):
-            shutil.rmtree(dropstat_path)
+        remove_tmp_dropstat_file()
         os.environ["PATH"] += os.pathsep + scripts_path
         os.environ["UTILITIES_UNIT_TESTING"] = "1"