# -*- coding: utf-8 -*-
#
# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.

"""Unit tests."""

import mock
import pandas
import pytest

from google.api_core import exceptions
from google.auth.credentials import AnonymousCredentials
from google.cloud import automl_v1beta1
from google.cloud.automl_v1beta1.proto import data_types_pb2

PROJECT = "project"
REGION = "region"
LOCATION_PATH = "projects/{}/locations/{}".format(PROJECT, REGION)


class TestTablesClient(object):
    def tables_client(
        self, client_attrs={}, prediction_client_attrs={}, gcs_c
lient_attrs={}
    ):
        client_mock = mock.Mock(**client_attrs)
        prediction_client_mock = mock.Mock(**prediction_client_attrs)
        gcs_client_mock = mock.Mock(**gcs_client_attrs)
        return automl_v1beta1.TablesClient(
            client=client_mock,
            prediction_
client=prediction_client_mock,
            gcs_client=gcs_client_mock,
            project=PROJECT,
            region=REGION,
        )

    def test_list_datasets_empty(self):
        client = self.tables_client(
            {
                "list_datasets.return_value": [],
                "location_path.return_value": LOCATION_PATH,
            },
            {},
        )
        ds = client.list_datasets()
        client.auto_ml_client.location_path.assert_called_with(PROJECT, REGION)
        client.auto_ml_client.list_datasets.assert_called_with(LOCATION_PATH)
        assert ds == []

    def test_list_datasets_not_empty(self):
        datasets = ["some_dataset"]
        client = self.tables_client(
            {
                "list_datasets.return_value": datasets,
                "location_path.return_value": LOCATION_PATH,
            },
            {},
        )
        ds = client.list_datasets()
        client.auto_ml_client.location_path.assert_called_with(PROJECT, REGION)
        client.auto_ml_client.list_datasets.assert_called_with(LOCATION_PATH)
        assert len(ds) == 1
        assert ds[0] == "some_dataset"

    def test_get_dataset_no_value(self):
        dataset_actual = "dataset"
        client = self.tables_client({}, {})
        with pytest.raises(ValueError):
            dataset = client.get_dataset()
        client.auto_ml_client.get_dataset.assert_not_called()

    def test_get_dataset_name(self):
        dataset_actual = "dataset"
        client = self.tables_client({"get_dataset.return_value": dataset_actual}, {})
        dataset = client.get_dataset(dataset_name="my_dataset")
        client.auto_ml_client.get_dataset.assert_called_with("my_dataset")
        assert dataset == dataset_actual

    def test_get_no_dataset(self):
        client = self.tables_client(
            {"get_dataset.side_effect": exceptions.NotFound("err")}, {}
        )
        with pytest.raises(exceptions.NotFound):
            client.get_dataset(dataset_name="my_dataset")
        client.auto_ml_client.get_dataset.assert_called_with("my_dataset")

    def test_get_dataset_from_empty_list(self):
        client = self.tables_client({"list_datasets.return_value": []}, {})
        with pytest.raises(exceptions.NotFound):
            client.get_dataset(dataset_display_name="my_dataset")

    def test_get_dataset_from_list_not_found(self):
        client = self.tables_client(
            {"list_datasets.return_value": [mock.Mock(display_name="not_it")]}, {}
        )
        with pytest.raises(exceptions.NotFound):
            client.get_dataset(dataset_display_name="my_dataset")

    def test_get_dataset_from_list(self):
        client = self.tables_client(
            {
                "list_datasets.return_value": [
                    mock.Mock(display_name="not_it"),
                    mock.Mock(display_name="my_dataset"),
                ]
            },
            {},
        )
        dataset = client.get_dataset(dataset_display_name="my_dataset")
        assert dataset.display_name == "my_dataset"

    def test_get_dataset_from_list_ambiguous(self):
        client = self.tables_client(
            {
                "list_datasets.return_value": [
                    mock.Mock(display_name="my_dataset"),
                    mock.Mock(display_name="not_my_dataset"),
                    mock.Mock(display_name="my_dataset"),
                ]
            },
            {},
        )
        with pytest.raises(ValueError):
            client.get_dataset(dataset_display_name="my_dataset")

    def test_create_dataset(self):
        client = self.tables_client(
            {
                "location_path.return_value": LOCATION_PATH,
                "create_dataset.return_value": mock.Mock(display_name="name"),
            },
            {},
        )
        metadata = {"metadata": "values"}
        dataset = client.create_dataset("name", metadata=metadata)
        client.auto_ml_client.location_path.assert_called_with(PROJECT, REGION)
        client.auto_ml_client.create_dataset.assert_called_with(
            LOCATION_PATH, {"display_name": "name", "tables_dataset_metadata": metadata}
        )
        assert dataset.display_name == "name"

    def test_delete_dataset(self):
        dataset = mock.Mock()
        dataset.configure_mock(name="name")
        client = self.tables_client({"delete_dataset.return_value": None}, {})
        client.delete_dataset(dataset=dataset)
        client.auto_ml_client.delete_dataset.assert_called_with("name")

    def test_delete_dataset_not_found(self):
        client = self.tables_client({"list_datasets.return_value": []}, {})
        client.delete_dataset(dataset_display_name="not_found")
        client.auto_ml_client.delete_dataset.assert_not_called()

    def test_delete_dataset_name(self):
        client = self.tables_client({"delete_dataset.return_value": None}, {})
        client.delete_dataset(dataset_name="name")
        client.auto_ml_client.delete_dataset.assert_called_with("name")

    def test_export_not_found(self):
        client = self.tables_client({"list_datasets.return_value": []}, {})
        with pytest.raises(exceptions.NotFound):
            client.export_data(dataset_display_name="name", gcs_input_uris="uri")

        client.auto_ml_client.export_data.assert_not_called()

    def test_export_gcs_uri(self):
        client = self.tables_client({"export_data.return_value": None}, {})
        client.export_data(dataset_name="name", gcs_output_uri_prefix="uri")
        client.auto_ml_client.export_data.assert_called_with(
            "name", {"gcs_destination": {"output_uri_prefix": "uri"}}
        )

    def test_export_bq_uri(self):
        client = self.tables_client({"export_data.return_value": None}, {})
        client.export_data(dataset_name="name", bigquery_output_uri="uri")
        client.auto_ml_client.export_data.assert_called_with(
            "name", {"bigquery_destination": {"output_uri": "uri"}}
        )

    def test_import_not_found(self):
        client = self.tables_client({"list_datasets.return_value": []}, {})
        with pytest.raises(exceptions.NotFound):
            client.import_data(dataset_display_name="name", gcs_input_uris="uri")

        client.auto_ml_client.import_data.assert_not_called()

    def test_import_pandas_dataframe(self):
        client = self.tables_client(
            gcs_client_attrs={
                "bucket_name": "my_bucket",
                "upload_pandas_dataframe.return_value": "uri",
            }
        )
        dataframe = pandas.DataFrame({})
        client.import_data(
            project=PROJECT,
            region=REGION,
            dataset_name="name",
            pandas_dataframe=dataframe,
        )
        client.gcs_client.ensure_bucket_exi