chore: tests

Signed-off-by: SdgJlbl <[email protected]>

substrafl/remote/operations.py

@@ -1,6 +1,7 @@
 """Dataclasses describing the operations
 to execute on the remote.
 """
+
 from dataclasses import dataclass
 from typing import Any
 from typing import List

substrafl/remote/register/register.py

@@ -53,7 +53,6 @@
 """
 
 DOCKERFILE_TEMPLATE = """{base_docker_image}
-
 # create a non-root user
 RUN addgroup --gid 1001 group
 RUN adduser --disabled-password --gecos "" --uid 1001 --gid 1001 --home /home/user user
@@ -129,7 +128,6 @@ def _check_python_version(python_major_minor: str) -> None:
 
 def _get_base_docker_image(python_major_minor: str, use_gpu: bool) -> str:
     """Get the base Docker image for the Dockerfile"""
-    _check_python_version(python_major_minor)
 
     if use_gpu:
         base_docker_image = _GPU_BASE_IMAGE.format(
@@ -154,6 +152,9 @@ def _create_dockerfile(install_libraries: bool, dependencies: Dependency, operat
     # Cloudpickle will crash if we don't deserialize with the same major.minor
     python_major_minor = ".".join(python_version().split(".")[:2])
 
+    # check that the Python version is supported
+    _check_python_version(python_major_minor)
+
     # Get the base Docker image
     base_docker_image = _get_base_docker_image(python_major_minor=python_major_minor, use_gpu=dependencies.use_gpu)
     # Build Substrafl, Substra and Substratools, and local dependencies wheels if necessary

substrafl/remote/serializers/__init__.py

@@ -1,6 +1,7 @@
 """
 Serializers to save the user code and wrap it in the Substra algo code.
 """
+
 from substrafl.remote.serializers.pickle_serializer import PickleSerializer
 from substrafl.remote.serializers.serializer import Serializer
 

substrafl/strategies/schemas.py

@@ -1,5 +1,6 @@
 """Schemas used in the strategies.
 """
+
 from enum import Enum
 from typing import List
 

tests/dependency/installable_library/setup.py

@@ -1,4 +1,5 @@
 """Packaging settings."""
+
 from setuptools import setup
 
 setup(

tests/dependency/installable_library2/setup.py

@@ -1,4 +1,5 @@
 """Packaging settings."""
+
 from setuptools import setup
 
 setup(

tests/installable_library/setup.py

@@ -1,4 +1,5 @@
 """Packaging settings."""
+
 from setuptools import find_packages
 from setuptools import setup
 

tests/remote/register/test_register.py

@@ -12,6 +12,7 @@
 from substrafl.remote.decorators import remote_data
 from substrafl.remote.register import register
 from substrafl.remote.register.register import _create_dockerfile
+from substrafl.remote.register.register import _get_base_docker_image
 
 
 class RemoteClass:
@@ -37,13 +38,42 @@ def test_check_python_version(version):
         register._check_python_version(version)
 
 
-@pytest.mark.parametrize("version", ["3.9", "3.10", "3.11"])
+@pytest.mark.parametrize("version", ["3.9", "3.10", "3.11", "3.12"])
 def test_check_python_version_valid(version):
     """Does not raise for supported versions"""
     register._check_python_version(version)
 
 
-def test_create_cpu_dockerfile(tmp_path, mocker, local_installable_module):
+def test_get_base_docker_image_cpu():
+    expected_dockerfile = """
+FROM python:3.12-slim
+
+# update image
+RUN apt-get update -y
+"""
+    assert expected_dockerfile == _get_base_docker_image("3.12", use_gpu=False)
+
+
+def test_get_base_docker_image_gpu():
+    expected_dockerfile = """
+FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
+
+# update image
+RUN apt-get update -y
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get install -y software-properties-common
+RUN add-apt-repository -y ppa:deadsnakes/ppa
+RUN apt-get -y upgrade
+
+# install python
+RUN apt-get install -y python3.11 python3.11-venv python3-pip
+
+"""
+    assert expected_dockerfile == _get_base_docker_image("3.11", use_gpu=True)
+
+
+def test_create_dockerfile(tmp_path, local_installable_module):
     python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
     substrafl_wheel = f"substrafl_internal/dist/substrafl-{substrafl.__version__}-py3-none-any.whl"
     substra_wheel = f"substrafl_internal/dist/substra-{substra.__version__}-py3-none-any.whl"

tests/settings.py

@@ -1,4 +1,5 @@
 """Global settings for all tests environment."""
+
 import functools
 from pathlib import Path
 from typing import List