# Copyright 2019 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 smexperiments import trial_component, api_types

import datetime
import pytest
import unittest.mock


@pytest.fixture
def sagemaker_boto_client():
    return unittest.mock.Mock()


def test_create(sagemaker_boto_client):
    sagemaker_boto_client.create_trial_component.return_value = {
        "TrialComponentArn": "bazz",
    }
    obj = trial_component.TrialComponent.create(
        trial_component_name="foo", display_name="bar", sagemaker_boto_client=sagemaker_boto_client
    )
    sagemaker_boto_client.create_trial_component.assert_called_with(TrialComponentName="foo", DisplayName="bar")
    assert "foo" == obj.trial_component_name
    assert "bar" == obj.display_name
    assert "bazz" == obj.trial_component_arn


def test_load(sagemaker_boto_client):
    now = datetime.datetime.now(datetime.timezone.utc)

    sagemaker_boto_client.describe_trial_component.return_value = {
        "TrialComponentArn": "A",
        "TrialComponentName": "B",
        "DisplayName": "C",
        "Status": {"PrimaryStatus": "InProgress", "Message": "D"},
        "Parameters": {"E": {"NumberValue": 1.0}, "F": {"StringValue": "G"}},
        "InputArtifacts": {"H": {"Value": "s3://foo/bar", "MediaType": "text/plain"}},
        "OutputArtifacts": {"I": {"Value": "s3://whizz/bang", "MediaType": "text/plain"}},
        "Metrics": [
            {
                "MetricName": "J",
                "Count": 1,
                "Min": 1.0,
                "Max": 2.0,
                "Avg": 3.0,
                "StdDev": 4.0,
                "SourceArn": "K",
                "Timestamp": now,
            }
        ],
    }
    obj = trial_component.TrialComponent.load(trial_component_name="foo", sagemaker_boto_client=sagemaker_boto_client)
    sagemaker_boto_client.describe_trial_component.assert_called_with(TrialComponentName="foo")
    assert "A" == obj.trial_component_arn
    assert "B" == obj.trial_component_name
    assert "C" == obj.display_name
    assert api_types.TrialComponentStatus(primary_status="InProgress", message="D") == obj.status
    assert {"E": 1.0, "F": "G"} == obj.parameters
    assert {"H": api_types.TrialComponentArtifact(value="s3://foo/bar", media_type="text/plain")}
    assert {"I": api_types.TrialComponentArtifact(value="s3://whizz/bang", media_type="text/plain")}
    assert [
        api_types.TrialComponentMetricSummary(
            metric_name="J", count=1, min=1.0, max=2.0, avg=3.0, std_dev=4.0, source_arn="K", timestamp=now
        )
    ]


def test_list(sagemaker_boto_client):
    start_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=1)
    end_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=2)
    creation_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=3)
    last_modified_time = datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=4)

    sagemaker_boto_client.list_trial_components.side_effect = [
        {
            "TrialComponentSummaries": [
                {
                    "TrialComponentName": "A" + str(i),
                    "TrialComponentArn": "B" + str(i),
                    "DisplayName": "C" + str(i),
                    "SourceArn": "D" + str(i),
                    "Status": {"PrimaryStatus": "InProgress", "Message": "E" + str(i)},
                    "StartTime": start_time + datetime.timedelta(hours=i),
                    "EndTime": end_time + datetime.timedelta(hours=i),
                    "CreationTime": creation_time + datetime.timedelta(hours=i),
                    "LastModifiedTime": last_modified_time + datetime.timedelta(hours=i),
                    "LastModifiedBy": {},
                }
                for i in range(10)
            ],
            "NextToken": "100",
        },
        {
            "TrialComponentSummaries": [
                {
                    "TrialComponentName": "A" + str(i),
                    "TrialComponentArn": "B" + str(i),
                    "DisplayName": "C" + str(i),
                    "SourceArn": "D" + str(i),
                    "Status": {"PrimaryStatus": "InProgress", "Message": "E" + str(i)},
                    "StartTime": start_time + datetime.timedelta(hours=i),
                    "EndTime": end_time + datetime.timedelta(hours=i),
                    "CreationTime": creation_time + datetime.timedelta(hours=i),
                    "LastModifiedTime": last_modified_time + datetime.timedelta(hours=i),
                    "LastModifiedBy": {},
                }
                for i in range(10, 20)
            ]
        },
    ]

    expected = [
        api_types.TrialComponentSummary(
            trial_component_name="A" + str(i),
            trial_component_arn="B" + str(i),
            display_name="C" + str(i),
            source_arn="D" + str(i),
            status=api_types.TrialComponentStatus(primary_status="InProgress", message="E" + str(i)),
            start_time=start_time + datetime.timedelta(hours=i),
            end_time=end_time + datetime.timedelta(hours=i),
            creation_time=creation_time + datetime.timedelta(hours=i),
            last_modified_time=last_modified_time + datetime.timedelta(hours=i),
            last_modified_by={},
        )
        for i in range(20)
    ]
    result = list(
        trial_component.TrialComponent.list(
            sagemaker_boto_client=sagemaker_boto_client,
            source_arn="foo",
            sort_by="CreationTime",
            sort_order="Ascending",
        )
    )

    assert expected == result
    expected_calls = [
        unittest.mock.call(SortBy="CreationTime", SortOrder="Ascending", SourceArn="foo"),
        unittest.mock.call(NextToken="100", SortBy="CreationTime", SortOrder="Ascending", SourceArn="foo"),
    ]
    assert expected_calls == sagemaker_boto_client.list_trial_components.mock_calls


def test_list_empty(sagemaker_boto_client):
    sagemaker_boto_client.list_trial_components.return_value = {"TrialComponentSummaries": []}
    assert [] == list(trial_component.TrialComponent.list(sagemaker_boto_client=sagemaker_boto_client))


def test_list_trial_components_call_args(sagemaker_boto_client):
    created_before = datetime.datetime(1999, 10, 12, 0, 0, 0)
    created_after = datetime.datetime(1990, 10, 12, 0, 0, 0)
    trial_name = "foo-trial"
    experiment_name = "foo-experiment"
    next_token = "thetoken"
    max_results = 99

    sagemaker_boto_client.list_trial_components.return_value = {}
    assert [] == list(
        trial_component.TrialComponent.list(
            sagemaker_boto_client=sagemaker_boto_client,
            trial_name=trial_name,
            experiment_name=experiment_name,
            created_before=created_before,
            created_after=created_after,
            next_token=next_token,
            max_results=max_results,
            sort_by="CreationTime",
            sort_order="Ascending",
        )
    )

    expected_calls = [
        unittest.mock.call(
            TrialName="foo-trial",
            ExperimentName="foo-experiment",
            CreatedBefore=created_before,
            CreatedAfter=created_after,
            SortBy="CreationTime",
            SortOrder="Ascending",
            NextToken="thetoken",
            MaxResults=99,
        )
    ]
    assert expected_calls == sagemaker_boto_client.list_trial_components.mock_calls


def test_save(sagemaker_boto_client):
    obj = trial_component.TrialComponent(sagemaker_boto_client, trial_component_name="foo", display_name="bar")
    sagemaker_boto_client.update_trial_component.return_value = {}
    obj.save()
    sagemaker_boto_client.update_trial_component.assert_called_with(TrialComponentName="foo", DisplayName="bar")


def test_delete(sagemaker_boto_client):
    obj = trial_component.TrialComponent(sagemaker_boto_client, trial_component_name="foo", display_name="bar")
    sagemaker_boto_client.delete_trial_component.return_value = {}
    obj.delete()
    sagemaker_boto_client.delete_trial_component.assert_called_with(TrialComponentName="foo")


def test_boto_ignore():
    obj = trial_component.TrialComponent(sagemaker_boto_client, trial_component_name="foo", display_name="bar")
    assert obj._boto_ignore() == ["ResponseMetadata", "CreatedBy"]

[(29, 'smexperiments.trial_component.TrialComponent.create', 'trial_component.TrialComponent.create', False, 'from smexperiments import trial_component, api_types\n'), (39, 'datetime.datetime.now', 'datetime.datetime.now', False, 'import datetime\n'), (62, 'smexperiments.trial_component.TrialComponent.load', 'trial_component.TrialComponent.load', False, 'from smexperiments import trial_component, api_types\n'), (160, 'datetime.datetime', 'datetime.datetime', False, 'import datetime\n'), (161, 'datetime.datetime', 'datetime.datetime', False, 'import datetime\n'), (198, 'smexperiments.trial_component.TrialComponent', 'trial_component.TrialComponent', False, 'from smexperiments import trial_component, api_types\n'), (205, 'smexperiments.trial_component.TrialComponent', 'trial_component.TrialComponent', False, 'from smexperiments import trial_component, api_types\n'), (212, 'smexperiments.trial_component.TrialComponent', 'trial_component.TrialComponent', False, 'from smexperiments import trial_component, api_types\n'), (67, 'smexperiments.api_types.TrialComponentStatus', 'api_types.TrialComponentStatus', False, 'from smexperiments import trial_component, api_types\n'), (69, 'smexperiments.api_types.TrialComponentArtifact', 'api_types.TrialComponentArtifact', False, 'from smexperiments import trial_component, api_types\n'), (70, 'smexperiments.api_types.TrialComponentArtifact', 'api_types.TrialComponentArtifact', False, 'from smexperiments import trial_component, api_types\n'), (72, 'smexperiments.api_types.TrialComponentMetricSummary', 'api_types.TrialComponentMetricSummary', False, 'from smexperiments import trial_component, api_types\n'), (79, 'datetime.datetime.now', 'datetime.datetime.now', False, 'import datetime\n'), (79, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (80, 'datetime.datetime.now', 'datetime.datetime.now', False, 'import datetime\n'), (80, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (81, 'datetime.datetime.now', 'datetime.datetime.now', False, 'import datetime\n'), (81, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (82, 'datetime.datetime.now', 'datetime.datetime.now', False, 'import datetime\n'), (82, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (138, 'smexperiments.trial_component.TrialComponent.list', 'trial_component.TrialComponent.list', False, 'from smexperiments import trial_component, api_types\n'), (156, 'smexperiments.trial_component.TrialComponent.list', 'trial_component.TrialComponent.list', False, 'from smexperiments import trial_component, api_types\n'), (169, 'smexperiments.trial_component.TrialComponent.list', 'trial_component.TrialComponent.list', False, 'from smexperiments import trial_component, api_types\n'), (129, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (130, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (131, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (132, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (93, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (94, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (95, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (96, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (111, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (112, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (113, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n'), (114, 'datetime.timedelta', 'datetime.timedelta', False, 'import datetime\n')]