# 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"]

[((910, 1044), 'smexperiments.trial_component.TrialComponent.create', 'trial_component.TrialComponent.create', ([], {'trial_component_name': '"""foo"""', 'display_name': '"""bar"""', 'sagemaker_boto_client': 'sagemaker_boto_client'}), "(trial_component_name='foo',\n    display_name='bar', sagemaker_boto_client=sagemaker_boto_client)\n", (947, 1044), False, 'from smexperiments import trial_component, api_types\n'), ((1345, 1389), 'datetime.datetime.now', 'datetime.datetime.now', (['datetime.timezone.utc'], {}), '(datetime.timezone.utc)\n', (1366, 1389), False, 'import datetime\n'), ((2201, 2313), 'smexperiments.trial_component.TrialComponent.load', 'trial_component.TrialComponent.load', ([], {'trial_component_name': '"""foo"""', 'sagemaker_boto_client': 'sagemaker_boto_client'}), "(trial_component_name='foo',\n    sagemaker_boto_client=sagemaker_boto_client)\n", (2236, 2313), False, 'from smexperiments import trial_component, api_types\n'), ((6783, 6823), 'datetime.datetime', 'datetime.datetime', (['(1999)', '(10)', '(12)', '(0)', '(0)', '(0)'], {}), '(1999, 10, 12, 0, 0, 0)\n', (6800, 6823), False, 'import datetime\n'), ((6844, 6884), 'datetime.datetime', 'datetime.datetime', (['(1990)', '(10)', '(12)', '(0)', '(0)', '(0)'], {}), '(1990, 10, 12, 0, 0, 0)\n', (6861, 6884), False, 'import datetime\n'), ((8013, 8119), 'smexperiments.trial_component.TrialComponent', 'trial_component.TrialComponent', (['sagemaker_boto_client'], {'trial_component_name': '"""foo"""', 'display_name': '"""bar"""'}), "(sagemaker_boto_client, trial_component_name=\n    'foo', display_name='bar')\n", (8043, 8119), False, 'from smexperiments import trial_component, api_types\n'), ((8362, 8468), 'smexperiments.trial_component.TrialComponent', 'trial_component.TrialComponent', (['sagemaker_boto_client'], {'trial_component_name': '"""foo"""', 'display_name': '"""bar"""'}), "(sagemaker_boto_client, trial_component_name=\n    'foo', display_name='bar')\n", (8392, 8468), False, 'from smexperiments import trial_component, api_types\n'), ((8678, 8784), 'smexperiments.trial_component.TrialComponent', 'trial_component.TrialComponent', (['sagemaker_boto_client'], {'trial_component_name': '"""foo"""', 'display_name': '"""bar"""'}), "(sagemaker_boto_client, trial_component_name=\n    'foo', display_name='bar')\n", (8708, 8784), False, 'from smexperiments import trial_component, api_types\n'), ((2537, 2609), 'smexperiments.api_types.TrialComponentStatus', 'api_types.TrialComponentStatus', ([], {'primary_status': '"""InProgress"""', 'message': '"""D"""'}), "(primary_status='InProgress', message='D')\n", (2567, 2609), False, 'from smexperiments import trial_component, api_types\n'), ((2691, 2770), 'smexperiments.api_types.TrialComponentArtifact', 'api_types.TrialComponentArtifact', ([], {'value': '"""s3://foo/bar"""', 'media_type': '"""text/plain"""'}), "(value='s3://foo/bar', media_type='text/plain')\n", (2723, 2770), False, 'from smexperiments import trial_component, api_types\n'), ((2789, 2876), 'smexperiments.api_types.TrialComponentArtifact', 'api_types.TrialComponentArtifact', ([], {'value': '"""s3://whizz/bang"""', 'media_type': '"""text/plain"""'}), "(value='s3://whizz/bang', media_type=\n    'text/plain')\n", (2821, 2876), False, 'from smexperiments import trial_component, api_types\n'), ((2894, 3032), 'smexperiments.api_types.TrialComponentMetricSummary', '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'}), "(metric_name='J', count=1, min=1.0,\n    max=2.0, avg=3.0, std_dev=4.0, source_arn='K', timestamp=now)\n", (2931, 3032), False, 'from smexperiments import trial_component, api_types\n'), ((3114, 3158), 'datetime.datetime.now', 'datetime.datetime.now', (['datetime.timezone.utc'], {}), '(datetime.timezone.utc)\n', (3135, 3158), False, 'import datetime\n'), ((3161, 3188), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': '(1)'}), '(hours=1)\n', (3179, 3188), False, 'import datetime\n'), ((3204, 3248), 'datetime.datetime.now', 'datetime.datetime.now', (['datetime.timezone.utc'], {}), '(datetime.timezone.utc)\n', (3225, 3248), False, 'import datetime\n'), ((3251, 3278), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': '(2)'}), '(hours=2)\n', (3269, 3278), False, 'import datetime\n'), ((3299, 3343), 'datetime.datetime.now', 'datetime.datetime.now', (['datetime.timezone.utc'], {}), '(datetime.timezone.utc)\n', (3320, 3343), False, 'import datetime\n'), ((3346, 3373), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': '(3)'}), '(hours=3)\n', (3364, 3373), False, 'import datetime\n'), ((3399, 3443), 'datetime.datetime.now', 'datetime.datetime.now', (['datetime.timezone.utc'], {}), '(datetime.timezone.utc)\n', (3420, 3443), False, 'import datetime\n'), ((3446, 3473), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': '(4)'}), '(hours=4)\n', (3464, 3473), False, 'import datetime\n'), ((5895, 6050), 'smexperiments.trial_component.TrialComponent.list', 'trial_component.TrialComponent.list', ([], {'sagemaker_boto_client': 'sagemaker_boto_client', 'source_arn': '"""foo"""', 'sort_by': '"""CreationTime"""', 'sort_order': '"""Ascending"""'}), "(sagemaker_boto_client=\n    sagemaker_boto_client, source_arn='foo', sort_by='CreationTime',\n    sort_order='Ascending')\n", (5930, 6050), False, 'from smexperiments import trial_component, api_types\n'), ((6613, 6698), 'smexperiments.trial_component.TrialComponent.list', 'trial_component.TrialComponent.list', ([], {'sagemaker_boto_client': 'sagemaker_boto_client'}), '(sagemaker_boto_client=sagemaker_boto_client\n    )\n', (6648, 6698), False, 'from smexperiments import trial_component, api_types\n'), ((7100, 7412), 'smexperiments.trial_component.TrialComponent.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"""'}), "(sagemaker_boto_client=\n    sagemaker_boto_client, trial_name=trial_name, experiment_name=\n    experiment_name, created_before=created_before, created_after=\n    created_after, next_token=next_token, max_results=max_results, sort_by=\n    'CreationTime', sort_order='Ascending')\n", (7135, 7412), False, 'from smexperiments import trial_component, api_types\n'), ((5550, 5577), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': 'i'}), '(hours=i)\n', (5568, 5577), False, 'import datetime\n'), ((5611, 5638), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': 'i'}), '(hours=i)\n', (5629, 5638), False, 'import datetime\n'), ((5682, 5709), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': 'i'}), '(hours=i)\n', (5700, 5709), False, 'import datetime\n'), ((5763, 5790), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': 'i'}), '(hours=i)\n', (5781, 5790), False, 'import datetime\n'), ((3949, 3976), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': 'i'}), '(hours=i)\n', (3967, 3976), False, 'import datetime\n'), ((4020, 4047), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': 'i'}), '(hours=i)\n', (4038, 4047), False, 'import datetime\n'), ((4101, 4128), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': 'i'}), '(hours=i)\n', (4119, 4128), False, 'import datetime\n'), ((4191, 4218), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': 'i'}), '(hours=i)\n', (4209, 4218), False, 'import datetime\n'), ((4783, 4810), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': 'i'}), '(hours=i)\n', (4801, 4810), False, 'import datetime\n'), ((4854, 4881), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': 'i'}), '(hours=i)\n', (4872, 4881), False, 'import datetime\n'), ((4935, 4962), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': 'i'}), '(hours=i)\n', (4953, 4962), False, 'import datetime\n'), ((5025, 5052), 'datetime.timedelta', 'datetime.timedelta', ([], {'hours': 'i'}), '(hours=i)\n', (5043, 5052), False, 'import datetime\n')]