fix: tests

Signed-off-by: Guilhem Barthes <[email protected]>

backend/builder/image_builder/image_builder.py

@@ -7,17 +7,17 @@
 from django.conf import settings
 
 import orchestrator
+from builder import docker
 from builder import exceptions
-from builder.docker import container_image_exists
 from builder.exceptions import BuildError
 from builder.exceptions import BuildRetryError
 from builder.kubernetes import get_pod_logs
 from builder.kubernetes import pod_exists
 from builder.kubernetes import watch_pod
 from builder.volumes import get_docker_cache_pvc_name
+from substrapp.compute_tasks import datastore as ds
 from substrapp.compute_tasks import utils
 from substrapp.compute_tasks.compute_pod import Label
-from substrapp.compute_tasks.datastore import Datastore
 from substrapp.compute_tasks.volumes import get_worker_subtuple_pvc_name
 from substrapp.docker_registry import USER_IMAGE_REPOSITORY
 from substrapp.kubernetes_utils import delete_pod
@@ -59,10 +59,10 @@ def build_image_if_missing(channel: str, function: orchestrator.Function) -> Non
     """
     Build the container image and the ImageEntryPoint entry if they don't exist already
     """
-    datastore = Datastore(channel=channel)
+    datastore = ds.Datastore(channel=channel)
     container_image_tag = utils.container_image_tag_from_function(function)
     with lock_resource("image-build", container_image_tag, ttl=IMAGE_BUILD_TIMEOUT, timeout=IMAGE_BUILD_TIMEOUT):
-        if container_image_exists(container_image_tag):
+        if docker.container_image_exists(container_image_tag):
             logger.info("Reusing existing image", image=container_image_tag)
         else:
             asset_content = datastore.get_function(function)
@@ -155,7 +155,7 @@ def _build_container_image(path: str, tag: str) -> None:
 
     except Exception as e:
         # In case of concurrent builds, it may fail. Check if the image exists.
-        if container_image_exists(tag):
+        if docker.container_image_exists(tag):
             logger.warning(
                 f"Build of container image {tag} failed, probably because it was done by a concurrent build",
                 exc_info=True,

backend/builder/tests/conftest.py

@@ -0,0 +1,9 @@
+import pytest
+
+import orchestrator
+import orchestrator.mock as orc_mock
+
+
+@pytest.fixture
+def function() -> orchestrator.Function:
+    return orc_mock.FunctionFactory()

backend/builder/tests/test_image_builder.py

@@ -31,14 +31,16 @@ def test_build_image_if_missing_image_already_exists(mocker: MockerFixture, func
     m_container_image_exists.assert_called_once_with(function_image_tag)
 
 
+@pytest.mark.django_db
 def test_build_image_if_missing_image_build_needed(mocker: MockerFixture, function: orchestrator.Function):
-    ds = mocker.Mock()
     m_container_image_exists = mocker.patch("builder.docker.container_image_exists", return_value=False)
+    m_datastore = mocker.patch("substrapp.compute_tasks.datastore.Datastore")
     m_build_function_image = mocker.patch("builder.image_builder.image_builder._build_function_image")
     function_image_tag = utils.container_image_tag_from_function(function)
 
-    image_builder.build_image_if_missing(datastore=ds, function=function)
+    image_builder.build_image_if_missing(channel="channel", function=function)
 
+    m_datastore.assert_called_once()
     m_container_image_exists.assert_called_once_with(function_image_tag)
     m_build_function_image.assert_called_once()
     assert m_build_function_image.call_args.args[1] == function

backend/substrapp/tasks/tasks_compute_task.py

@@ -29,6 +29,7 @@
 from substrapp.clients import organization as organization_client
 from substrapp.compute_tasks import compute_task as task_utils
 from substrapp.compute_tasks import errors as compute_task_errors
+from substrapp.compute_tasks import image_builder
 from substrapp.compute_tasks.asset_buffer import add_assets_to_taskdir
 from substrapp.compute_tasks.asset_buffer import add_task_assets_to_buffer
 from substrapp.compute_tasks.asset_buffer import clear_assets_buffer
@@ -45,8 +46,6 @@
 from substrapp.compute_tasks.directories import restore_dir
 from substrapp.compute_tasks.directories import teardown_task_dirs
 from substrapp.compute_tasks.execute import execute_compute_task
-from substrapp.compute_tasks.image_builder import load_remote_function_image
-from substrapp.compute_tasks.image_builder import wait_for_image_built
 from substrapp.compute_tasks.lock import MAX_TASK_DURATION
 from substrapp.compute_tasks.lock import acquire_compute_plan_lock
 from substrapp.compute_tasks.outputs import OutputSaver
@@ -257,10 +256,10 @@ def _run(
         # start build_image timer
         timer.start()
 
-        wait_for_image_built(ctx.function, channel_name)
+        image_builder.wait_for_image_built(ctx.function, channel_name)
 
         if get_owner() != ctx.function.owner:
-            load_remote_function_image(ctx.function, channel_name)
+            image_builder.load_remote_function_image(ctx.function, channel_name)
 
         # stop build_image timer
         _create_task_profiling_step(channel_name, task.key, ComputeTaskSteps.BUILD_IMAGE, timer.stop())

backend/substrapp/tests/tasks/test_compute_task.py

@@ -32,7 +32,8 @@ def test_compute_task_exception(mocker: MockerFixture):
     mock_init_task_dirs = mocker.patch("substrapp.tasks.tasks_compute_task.init_task_dirs")
     mock_add_asset_to_buffer = mocker.patch("substrapp.tasks.tasks_compute_task.add_task_assets_to_buffer")
     mock_add_asset_to_task_dir = mocker.patch("substrapp.tasks.tasks_compute_task.add_assets_to_taskdir")
-    mock_build_image_if_missing = mocker.patch("builder.image_builder.image_builder.build_image_if_missing")
+    mock_load_remote_function_image = mocker.patch("substrapp.compute_tasks.image_builder.load_remote_function_image")
+    mock_wait_for_image_built = mocker.patch("substrapp.compute_tasks.image_builder.wait_for_image_built")
     mock_execute_compute_task = mocker.patch("substrapp.tasks.tasks_compute_task.execute_compute_task")
     saver = mocker.MagicMock()
     mock_output_saver = mocker.patch("substrapp.tasks.tasks_compute_task.OutputSaver", return_value=saver)
@@ -67,7 +68,8 @@ class FakeDirectories:
     mock_init_task_dirs.assert_called_once()
     mock_add_asset_to_buffer.assert_called_once()
     mock_add_asset_to_task_dir.assert_called_once()
-    mock_build_image_if_missing.assert_called_once()
+    mock_load_remote_function_image.assert_called_once()
+    mock_wait_for_image_built.assert_called_once()
     mock_execute_compute_task.assert_called_once()
     saver.save_outputs.assert_called_once()
     mock_output_saver.assert_called_once()
@@ -136,7 +138,8 @@ def test_celery_retry(mocker: MockerFixture):
     mocker.patch("substrapp.tasks.tasks_compute_task.add_task_assets_to_buffer")
     mocker.patch("substrapp.tasks.tasks_compute_task.add_assets_to_taskdir")
     mocker.patch("substrapp.tasks.tasks_compute_task.restore_dir")
-    mocker.patch("builder.image_builder.image_builder.build_image_if_missing")
+    mocker.patch("substrapp.compute_tasks.image_builder.load_remote_function_image")
+    mocker.patch("substrapp.compute_tasks.image_builder.wait_for_image_built")
     mock_execute_compute_task = mocker.patch("substrapp.tasks.tasks_compute_task.execute_compute_task")
     mocker.patch("substrapp.tasks.tasks_compute_task.teardown_task_dirs")
     mock_retry = mocker.patch("substrapp.tasks.tasks_compute_task.ComputeTask.retry")