From b07ddc61b8834b83d64c4b3b7fca5ee06c132ce4 Mon Sep 17 00:00:00 2001
From: bingwang-ms <66248323+bingwang-ms@users.noreply.github.com>
Date: Wed, 9 Oct 2024 17:58:29 -0700
Subject: [PATCH] Workaround python_sai_thrift path issue (#14886)

---
 tests/test_pretest.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tests/test_pretest.py b/tests/test_pretest.py
index 061b3b7d955..2f406a62c80 100644
--- a/tests/test_pretest.py
+++ b/tests/test_pretest.py
@@ -329,6 +329,13 @@ def test_update_saithrift_ptf(request, ptfhost):
     if result["failed"] or "OK" not in result["msg"]:
         pytest.skip("Download failed/error while installing python saithrift package")
     ptfhost.shell("dpkg -i {}".format(os.path.join("/root", pkg_name)))
+    # In 202405 branch, the switch_sai_thrift package is inside saithrift-0.9-py3.11.egg
+    # We need to move it out to the correct location
+    PY_PATH = "/usr/lib/python3/dist-packages/"
+    SRC_PATH = PY_PATH + "saithrift-0.9-py3.11.egg/switch_sai_thrift"
+    DST_PATH = PY_PATH + "switch_sai_thrift"
+    if ptfhost.stat(path=SRC_PATH)['stat']['exists'] and not ptfhost.stat(path=DST_PATH)['stat']['exists']:
+        ptfhost.copy(src=SRC_PATH, dest=DST_PATH, remote_src=True)
     logging.info("Python saithrift package installed successfully")