def fetch_dlc_images_for_test_jobs()

in src/utils.py [0:0]

• 32 lines of code
• 14 McCabe index (conditional complexity)

def fetch_dlc_images_for_test_jobs(images):
    """
    use the JobParamters.run_test_types values to pass on image ecr urls to each test type.

    Args:
        images: list

    Returns:
        dictionary
    """
    DLC_IMAGES = {"sagemaker": [], "ecs": [], "eks": [], "ec2": [], "sanity": []}

    for docker_image in images:
        if docker_image.build_status == constants.SUCCESS:
            # Run sanity tests on the all images built
            DLC_IMAGES["sanity"].append(docker_image.ecr_url)
            image_job_type = docker_image.info.get("image_type")
            image_device_type = docker_image.info.get("device_type")
            image_python_version = docker_image.info.get("python_version")
            image_tag = f"{image_job_type}_{image_device_type}_{image_python_version}"
            # when image_run_test_types has key all values can be (all , ecs, eks, ec2, sagemaker)
            if constants.ALL in JobParameters.image_run_test_types.keys():
                run_tests = JobParameters.image_run_test_types.get(constants.ALL)
                run_tests = (
                    constants.ALL_TESTS if constants.ALL in run_tests else run_tests
                )
                for test in run_tests:
                    DLC_IMAGES[test].append(docker_image.ecr_url)
            # when key is training or inference values can be  (ecs, eks, ec2, sagemaker)
            if image_job_type in JobParameters.image_run_test_types.keys():
                run_tests = JobParameters.image_run_test_types.get(image_job_type)
                for test in run_tests:
                    DLC_IMAGES[test].append(docker_image.ecr_url)
            # when key is image_tag (training-cpu-py3) values can be (ecs, eks, ec2, sagemaker)
            if image_tag in JobParameters.image_run_test_types.keys():
                run_tests = JobParameters.image_run_test_types.get(image_tag)
                run_tests = (
                    constants.ALL_TESTS if constants.ALL in run_tests else run_tests
                )
                for test in run_tests:
                    DLC_IMAGES[test].append(docker_image.ecr_url)

    for test_type in DLC_IMAGES.keys():
        test_images = DLC_IMAGES[test_type]
        if test_images:
            DLC_IMAGES[test_type] = list(set(test_images))
    return DLC_IMAGES