diff --git a/python_files/tests/pytestadapter/helpers.py b/python_files/tests/pytestadapter/helpers.py
index dd69379a45b9..c768749c92fc 100644
--- a/python_files/tests/pytestadapter/helpers.py
+++ b/python_files/tests/pytestadapter/helpers.py
@@ -11,10 +11,14 @@
 import sys
 import threading
 import uuid
-from typing import Any, Dict, List, Optional, Tuple, TypedDict
+from typing import Any, Dict, List, Optional, Tuple
 
+script_dir = pathlib.Path(__file__).parent.parent.parent
+sys.path.append(os.fspath(script_dir))
+sys.path.append(os.fspath(script_dir / "lib" / "python"))
 
 TEST_DATA_PATH = pathlib.Path(__file__).parent / ".data"
+from typing_extensions import TypedDict  # noqa: E402
 
 
 def get_absolute_test_id(test_id: str, testPath: pathlib.Path) -> str:
diff --git a/python_files/vscode_pytest/__init__.py b/python_files/vscode_pytest/__init__.py
index 54bed66deae0..6c820195669b 100644
--- a/python_files/vscode_pytest/__init__.py
+++ b/python_files/vscode_pytest/__init__.py
@@ -10,9 +10,13 @@
 
 import pytest
 
-from typing import Any, Dict, List, Optional, Union, Literal, TypedDict  # noqa: E402
+from typing import Any, Dict, List, Optional, Union
 
+script_dir = pathlib.Path(__file__).parent.parent
+sys.path.append(os.fspath(script_dir))
+sys.path.append(os.fspath(script_dir / "lib" / "python"))
 
+from typing_extensions import Literal, TypedDict  # noqa: E402
 from testing_tools import socket_manager  # noqa: E402
 
 DEFAULT_PORT = 45454