# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
from __future__ import absolute_import

import pytest
from mock import Mock
from sagemaker.mxnet import MXNet


SCRIPT = "resnet_cifar_10.py"
TIMESTAMP = "2017-11-06-14:14:15.673"
TIME = 1510006209.073025
BUCKET_NAME = "mybucket"
INSTANCE_COUNT = 1
INSTANCE_TYPE_GPU = "ml.p2.xlarge"
INSTANCE_TYPE_CPU = "ml.m4.xlarge"
CPU_IMAGE_NAME = "sagemaker-mxnet-py2-cpu"
GPU_IMAGE_NAME = "sagemaker-mxnet-py2-gpu"
REGION = "us-west-2"
IMAGE_URI_FORMAT_STRING = "520713654638.dkr.ecr.{}.amazonaws.com/{}:{}-{}-{}"
REGION = "us-west-2"
ROLE = "SagemakerRole"
SOURCE_DIR = "s3://fefergerger"


@pytest.fixture()
def sagemaker_session():
    boto_mock = Mock(name="boto_session", region_name=REGION)
    ims = Mock(
        name="sagemaker_session",
        boto_session=boto_mock,
        config=None,
        local_mode=False,
        region_name=REGION,
    )

    ims.default_bucket = Mock(name="default_bucket", return_value=BUCKET_NAME)
    ims.expand_role = Mock(name="expand_role", return_value=ROLE)
    ims.sagemaker_client.describe_training_job = Mock(
        return_value={"ModelArtifacts": {"S3ModelArtifacts": "s3://m/m.tar.gz"}}
    )
    return ims


# Test that we pass all necessary fields from estimator to the session when we call deploy
def test_deploy(sagemaker_session, tf_version):
    estimator = MXNet(
        entry_point=SCRIPT,
        source_dir=SOURCE_DIR,
        role=ROLE,
        framework_version=tf_version,
        train_instance_count=2,
        train_instance_type=INSTANCE_TYPE_GPU,
        sagemaker_session=sagemaker_session,
        base_job_name="test-cifar",
    )

    estimator.fit("s3://mybucket/train")
    print("job succeeded: {}".format(estimator.latest_training_job.name))

    estimator.deploy(initial_instance_count=1, instance_type=INSTANCE_TYPE_CPU)
    image = IMAGE_URI_FORMAT_STRING.format(REGION, CPU_IMAGE_NAME, tf_version, "cpu", "py2")
    sagemaker_session.create_model.assert_called_with(
        estimator._current_job_name,
        ROLE,
        {
            "Environment": {
                "SAGEMAKER_CONTAINER_LOG_LEVEL": "20",
                "SAGEMAKER_SUBMIT_DIRECTORY": SOURCE_DIR,
                "SAGEMAKER_REGION": REGION,
                "SAGEMAKER_PROGRAM": SCRIPT,
            },
            "Image": image,
            "ModelDataUrl": "s3://m/m.tar.gz",
        },
    )