Fix worker pool integration tests

Revert to use a different dockerfile for each test

tests/integration/container_workload/pool_image_test.py

@@ -12,7 +12,6 @@
 from syft.client.domain_client import DomainClient
 from syft.custom_worker.config import DockerWorkerConfig
 from syft.custom_worker.config import PrebuiltWorkerConfig
-from syft.custom_worker.config import WorkerConfig
 from syft.service.request.request import Request
 from syft.service.response import SyftSuccess
 from syft.service.worker.worker_image import SyftWorkerImage
@@ -63,15 +62,14 @@ def external_registry_uid(domain_1_port: int) -> UID:
 
 PREBUILT_WORKER_CONFIG = PrebuiltWorkerConfig(tag=PREBUILT_IMAGE_TAG)
 
-CUSTOM_DOCKERFILE = f"""
-FROM {registry}/{repo}:{tag}
 
-RUN pip install recordlinkage
-"""
-
-CUSTOM_IMAGE_TAG = "openmined/custom-worker-recordlinkage"
-
-CUSTOM_DOCKER_WORKER_CONFIG = DockerWorkerConfig(dockerfile=CUSTOM_DOCKERFILE)
+def make_docker_config_test_case(pkg: str) -> tuple[str, str]:
+ return (
+ DockerWorkerConfig(
+ dockerfile=(f"FROM {registry}/{repo}:{tag}\nRUN pip install {pkg}\n")
+ ),
+ f"openmined/custom-worker-{pkg}:latest",
+ )
 
 
 @pytest.mark.container_workload
@@ -80,20 +78,19 @@ def test_image_build(domain_1_port: int, external_registry_uid: UID) -> None:
  port=domain_1_port, email="[email protected]", password="changethis"
  )
 
+ docker_config, docker_tag = make_docker_config_test_case("recordlinkage")
+
  submit_result = domain_client.api.services.worker_image.submit_container_image(
- worker_config=CUSTOM_DOCKER_WORKER_CONFIG
+ worker_config=docker_config
  )
  assert isinstance(submit_result, SyftSuccess)
  assert len(domain_client.images.get_all()) == 2
 
  # Validate if we can get the worker image object from its config
- workerimage = domain_client.api.services.worker_image.get_by_config(
- CUSTOM_DOCKER_WORKER_CONFIG
- )
+ workerimage = domain_client.api.services.worker_image.get_by_config(docker_config)
  assert not isinstance(workerimage, sy.SyftError)
 
  # Build docker image
- docker_tag = f"{CUSTOM_IMAGE_TAG}:{sy.UID().short()}"
  docker_build_result = domain_client.api.services.worker_image.build(
  image_uid=workerimage.id,
  tag=docker_tag,
@@ -112,18 +109,21 @@ def test_image_build(domain_1_port: int, external_registry_uid: UID) -> None:
 
 
 @pytest.mark.container_workload
-@pytest.mark.parametrize(
- "worker_config", [PREBUILT_WORKER_CONFIG, CUSTOM_DOCKER_WORKER_CONFIG]
-)
+@pytest.mark.parametrize("prebuilt", [True, False])
 def test_pool_launch(
- domain_1_port: int, external_registry_uid: UID, worker_config: WorkerConfig
+ domain_1_port: int, external_registry_uid: UID, prebuilt: bool
 ) -> None:
  domain_client: DomainClient = sy.login(
  port=domain_1_port, email="[email protected]", password="changethis"
  )
  assert len(domain_client.worker_pools.get_all()) == 1
 
  # Submit Worker Image
+ worker_config, docker_tag = (
+ (PREBUILT_WORKER_CONFIG, PREBUILT_IMAGE_TAG)
+ if prebuilt
+ else make_docker_config_test_case("opendp")
+ )
  submit_result = domain_client.api.services.worker_image.submit_container_image(
  worker_config=worker_config
  )
@@ -137,21 +137,20 @@ def test_pool_launch(
  assert not worker_image.is_built
 
  # Build docker image
- docker_tag = f"{CUSTOM_IMAGE_TAG}:{sy.UID().short()}"
  docker_build_result = domain_client.api.services.worker_image.build(
  image_uid=worker_image.id,
  tag=docker_tag,
  pull=False,
  )
  assert isinstance(docker_build_result, SyftSuccess)
 
- # Push Image to External registry
- push_result = domain_client.api.services.worker_image.push(
- worker_image.id,
- username=external_registry_username,
- password=external_registry_password,
- )
- assert isinstance(push_result, sy.SyftSuccess), str(push_result)
+  # Push Image to External registry
+  push_result = domain_client.api.services.worker_image.push(
+  worker_image.id,
+  username=external_registry_username,
+  password=external_registry_password,
+  )
+  assert isinstance(push_result, sy.SyftSuccess), str(push_result)
 
  # Launch a worker pool
  worker_pool_name = "custom-worker-pool-opendp"
@@ -198,11 +197,9 @@ def test_pool_launch(
 
 
 @pytest.mark.container_workload
-@pytest.mark.parametrize(
- "worker_config", [PREBUILT_WORKER_CONFIG, CUSTOM_DOCKER_WORKER_CONFIG]
-)
+@pytest.mark.parametrize("prebuilt", [True, False])
 def test_pool_image_creation_job_requests(
- domain_1_port: int, external_registry_uid: UID, worker_config: WorkerConfig
+ domain_1_port: int, external_registry_uid: UID, prebuilt: bool
 ) -> None:
  """
  Test register ds client, ds requests to create an image and pool creation,
@@ -226,11 +223,17 @@ def test_pool_image_creation_job_requests(
  ds_client = sy.login(email=ds_email, password="secret_pw", port=domain_1_port)
 
  # the DS makes a request to create an image and a pool based on the image
+ worker_config, docker_tag = (
+ (PREBUILT_WORKER_CONFIG, PREBUILT_IMAGE_TAG)
+ if prebuilt
+ else make_docker_config_test_case("numpy")
+ )
+
  worker_pool_name = "custom_worker_pool"
  request = ds_client.api.services.worker_pool.create_image_and_pool_request(
  pool_name=worker_pool_name,
  num_workers=1,
- tag=f"{CUSTOM_IMAGE_TAG}:{sy.UID().short()}",
+ tag=docker_tag,
  config=worker_config,
  reason="I want to do some more cool data science with PySyft",
  pull_image=False,