def test_byo_model_config(sagemaker_session):
    byo_model = model.Model(
        model_data="{{ model_data }}",
        image="{{ image }}",
        role="{{ role }}",
        env={"{{ key }}": "{{ value }}"},
        name="model",
        sagemaker_session=sagemaker_session,
    )

    config = airflow.model_config(instance_type="ml.c4.xlarge",
                                  model=byo_model)
    expected_config = {
        "ModelName": "model",
        "PrimaryContainer": {
            "Image": "{{ image }}",
            "Environment": {
                "{{ key }}": "{{ value }}"
            },
            "ModelDataUrl": "{{ model_data }}",
        },
        "ExecutionRoleArn": "{{ role }}",
    }

    assert config == expected_config
def test_byo_model_config(sagemaker_session):
    byo_model = model.Model(
        model_data="{{ model_data }}",
        image="{{ image }}",
        role="{{ role }}",
        env={"{{ key }}": "{{ value }}"},
        name='model',
        sagemaker_session=sagemaker_session)

    config = airflow.model_config(instance_type='ml.c4.xlarge', model=byo_model)
    expected_config = {
        'ModelName': 'model',
        'PrimaryContainer': {
            'Image': '{{ image }}',
            'Environment': {'{{ key }}': '{{ value }}'},
            'ModelDataUrl': '{{ model_data }}'
        },
        'ExecutionRoleArn': '{{ role }}'
    }

    assert config == expected_config