SNOW-1418543 make local testing thread safe

src/snowflake/snowpark/mock/_functions.py

@@ -10,6 +10,7 @@
 import operator
 import re
 import string
+import threading
 from decimal import Decimal
 from functools import partial, reduce
 from numbers import Real
@@ -130,14 +131,17 @@ def __call__(self, *args, input_data=None, row_number=None, **kwargs):
 
 class MockedFunctionRegistry:
     _instance = None
+    _lock_init = threading.Lock()
 
     def __init__(self) -> None:
         self._registry = dict()
+        self._lock = threading.RLock()
 
     @classmethod
     def get_or_create(cls) -> "MockedFunctionRegistry":
-        if cls._instance is None:
-            cls._instance = MockedFunctionRegistry()
+        with cls._lock_init:
+            if cls._instance is None:
+                cls._instance = MockedFunctionRegistry()
         return cls._instance
 
     def get_function(
@@ -151,10 +155,11 @@ def get_function(
             distinct = func.is_distinct
         func_name = func_name.lower()
 
-        if func_name not in self._registry:
-            return None
+        with self._lock:
+            if func_name not in self._registry:
+                return None
 
-        function = self._registry[func_name]
+            function = self._registry[func_name]
 
         return function.distinct if distinct else function
 
@@ -169,7 +174,8 @@ def register(
             snowpark_func if isinstance(snowpark_func, str) else snowpark_func.__name__
         )
         mocked_function = MockedFunction(name, func_implementation, *args, **kwargs)
-        self._registry[name] = mocked_function
+        with self._lock:
+            self._registry[name] = mocked_function
         return mocked_function
 
     def unregister(
@@ -180,8 +186,9 @@ def unregister(
             snowpark_func if isinstance(snowpark_func, str) else snowpark_func.__name__
         )
 
-        if name in self._registry:
-            del self._registry[name]
+        with self._lock:
+            if name in self._registry:
+                del self._registry[name]
 
 
 class LocalTimezone: