test fixes

Author: sfc-gh-bkogan

src/snowflake/snowpark/context.py

@@ -53,12 +53,13 @@
 
 # The fully qualified name of the Anaconda shared repository (conda channel).
 _ANACONDA_SHARED_REPOSITORY = "snowflake.snowpark.anaconda_shared_repository"
-# In case of failures or the current default artifact repository is unset, we fallback to this
-# TODO: undo this when the system function NULL behavior starts working again
+# The fully qualified name of the PyPI shared repository (pypi channel).
+_PYPI_SHARED_REPOSITORY = "snowflake.snowpark.pypi_shared_repository"
+# In case of failures and for routing to the right session package store, we use this
 _DEFAULT_ARTIFACT_REPOSITORY = (
     _ANACONDA_SHARED_REPOSITORY
     if sys.version_info <= (3, 13)
-    else "snowflake.snowpark.pypi_shared_repository"
+    else _PYPI_SHARED_REPOSITORY
 )
 
 

tests/integ/test_session.py

@@ -42,7 +42,7 @@
     _get_active_session,
     _get_active_sessions,
 )
-from snowflake.snowpark.context import _ANACONDA_SHARED_REPOSITORY
+from snowflake.snowpark.context import _DEFAULT_ARTIFACT_REPOSITORY
 from tests.utils import IS_IN_STORED_PROC, IS_IN_STORED_PROC_LOCALFS, TestFiles, Utils
 
 
@@ -1115,5 +1115,5 @@ def mock_session_parameters(param: str, quoted: bool = True):
         mock_session_parameters,
     ), caplog.at_level(logging.WARNING):
         result = session._get_default_artifact_repository()
-        assert result == _ANACONDA_SHARED_REPOSITORY
+        assert result == _DEFAULT_ARTIFACT_REPOSITORY
         assert caplog.text.count("Error getting default artifact repository") == 0

tests/unit/test_session.py

@@ -29,7 +29,10 @@
     SnowparkSessionException,
 )
 from snowflake.snowpark.session import _PYTHON_SNOWPARK_USE_SCOPED_TEMP_OBJECTS_STRING
-from snowflake.snowpark.context import _ANACONDA_SHARED_REPOSITORY
+from snowflake.snowpark.context import (
+    _ANACONDA_SHARED_REPOSITORY,
+    _DEFAULT_ARTIFACT_REPOSITORY,
+)
 from snowflake.snowpark.types import StructField, StructType
 
 
@@ -351,10 +354,10 @@ def assert_packages(packages):
 def test_resolve_packages_optional_artifact_repository(mock_server_connection):
     session = Session(mock_server_connection)
     session._get_default_artifact_repository = MagicMock(
-        return_value="snowflake.snowpark.pypi_shared_repository"
+        return_value="test_artifact_repo"
     )
     session._artifact_repository_packages = {
-        "snowflake.snowpark.pypi_shared_repository": {
+        "test_artifact_repo": {
             "numpy": "numpy==1.0.0",
         }
     }
@@ -762,6 +765,7 @@ def test_get_default_artifact_repository():
 
         assert mocked_run_query.call_count == 1
 
+    # TODO: remove this test when the system function NULL behavior starts working again?
     with mock.patch.object(
         session, "_run_query", return_value=[[None]]
     ) as mocked_run_query, mock.patch.object(
@@ -770,7 +774,7 @@ def test_get_default_artifact_repository():
         session, "get_current_schema", return_value="SCHEMA2"
     ):
         result = session._get_default_artifact_repository()
-        assert result == _ANACONDA_SHARED_REPOSITORY
+        assert result == _DEFAULT_ARTIFACT_REPOSITORY
 
         assert mocked_run_query.call_count == 1
 
@@ -799,6 +803,6 @@ def test_infer_is_return_table_uses_internal_describe():
         session, "get_current_schema", return_value="SCHEMA1"
     ):
         result = session._get_default_artifact_repository()
-        assert result == _ANACONDA_SHARED_REPOSITORY
+        assert result == _DEFAULT_ARTIFACT_REPOSITORY
 
         assert mocked_run_query.call_count == 1

tests/unit/test_stored_procedure.py

@@ -149,9 +149,10 @@ def mock_callback(extension_function_properties):
             == f"{sys.version_info[0]}.{sys.version_info[1]}"
         )
         assert callableProperties.all_imports == ""
+        # TODO: for python 3.14, it will use PyPI so we should also check for cloudpickle
         assert (
-            callableProperties.all_packages
-            == f"'snowflake-snowpark-python=={'.'.join(map(str, VERSION))}'"
+            f"'snowflake-snowpark-python=={'.'.join(map(str, VERSION))}'"
+            in callableProperties.all_packages
         )
         assert callableProperties.external_access_integrations is None
         assert callableProperties.secrets is None

tests/unit/test_udf_utils.py

@@ -149,7 +149,7 @@ def test_resolve_imports_and_packages_in_sandbox():
     assert handler is None
     assert inline_code is None
     assert all_imports == ""
-    assert all_packages == ",".join([f"'{package}'" for package in packages])
+    assert all(f"'{package}'" in all_packages for package in packages)
     assert upload_file_stage_location is None
     assert not custom_python_runtime_version_allowed
 
@@ -208,7 +208,7 @@ def test_resolve_imports_and_packages_imports_as_str(tmp_path_factory):
         assert handler is None
         assert inline_code is None
         assert all_imports == ""
-        assert all_packages == ",".join([f"'{package}'" for package in packages])
+        assert all(f"'{package}'" in all_packages for package in packages)
         assert upload_file_stage_location is None
         assert not custom_python_runtime_version_allowed