text: # Copyright 2015 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
  #
  #     http://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.
  
  import mock
  
  from .helpers import _Base
  from .helpers import _make_client
  from .helpers import _make_connection
  
  
  class TestCopyJobConfig(_Base):
      JOB_TYPE = "copy"
  
      @staticmethod
      def _get_target_class():
          from google.cloud.bigquery.job import CopyJobConfig
  
          return CopyJobConfig
  
      def test_ctor_defaults(self):
          from google.cloud.bigquery.job import OperationType
  
          config = self._make_one()
  
          assert config.create_disposition is None
          assert config.write_disposition is None
          assert config.destination_encryption_configuration is None
          assert config.operation_type == OperationType.OPERATION_TYPE_UNSPECIFIED
  
      def test_ctor_w_properties(self):
          from google.cloud.bigquery.job import CreateDisposition
          from google.cloud.bigquery.job import OperationType
          from google.cloud.bigquery.job import WriteDisposition
  
          create_disposition = CreateDisposition.CREATE_NEVER
          write_disposition = WriteDisposition.WRITE_TRUNCATE
          snapshot_operation = OperationType.SNAPSHOT
  
          config = self._get_target_class()(
              create_disposition=create_disposition,
              write_disposition=write_disposition,
              operation_type=snapshot_operation,
          )
  
          self.assertEqual(config.create_disposition, create_disposition)
          self.assertEqual(config.write_disposition, write_disposition)
          self.assertEqual(config.operation_type, snapshot_operation)
  
      def test_to_api_repr_with_encryption(self):
          from google.cloud.bigquery.encryption_configuration import (
              EncryptionConfiguration,
          )
  
          config = self._make_one()
          config.destination_encryption_configuration = EncryptionConfiguration(
              kms_key_name=self.KMS_KEY_NAME
          )
          resource = config.to_api_repr()
          self.assertEqual(
              resource,
              {
                  "copy": {
                      "destinationEncryptionConfiguration": {
                          "kmsKeyName": self.KMS_KEY_NAME
                      }
                  }
              },
          )
  
      def test_to_api_repr_with_encryption_none(self):
          config = self._make_one()
          config.destination_encryption_configuration = None
          resource = config.to_api_repr()
          self.assertEqual(
              resource, {"copy": {"destinationEncryptionConfiguration": None}}
          )
  
      def test_operation_type_setting_none(self):
          from google.cloud.bigquery.job import OperationType
  
          config = self._make_one(operation_type=OperationType.SNAPSHOT)
  
          # Setting it to None is the same as setting it to OPERATION_TYPE_UNSPECIFIED.
          config.operation_type = None
          assert config.operation_type == OperationType.OPERATION_TYPE_UNSPECIFIED
  
      def test_operation_type_setting_non_none(self):
          from google.cloud.bigquery.job import OperationType
  
          config = self._make_one(operation_type=None)
          config.operation_type = OperationType.RESTORE
          assert config.operation_type == OperationType.RESTORE
  
  
  class TestCopyJob(_Base):
      JOB_TYPE = "copy"
      SOURCE_TABLE = "source_table"
      DESTINATION_TABLE = "destination_table"
  
      @staticmethod
      def _get_target_class():
          from google.cloud.bigquery.job import CopyJob
  
          return CopyJob
  
      def _make_resource(self, started=False, ended=False):
          resource = super(TestCopyJob, self)._make_resource(started, ended)
          config = resource["configuration"]["copy"]
          config["sourceTables"] = [
              {
                  "projectId": self.PROJECT,
                  "datasetId": self.DS_ID,
                  "tableId": self.SOURCE_TABLE,
              }
          ]
          config["destinationTable"] = {
              "projectId": self.PROJECT,
              "datasetId": self.DS_ID,
              "tableId": self.DESTINATION_TABLE,
          }
  
          return resource
  
      def _verifyResourceProperties(self, job, resource):
          self._verifyReadonlyResourceProperties(job, resource)
  
          config = resource.get("configuration", {}).get("copy")
  
          table_ref = config["destinationTable"]
          self.assertEqual(job.destination.project, table_ref["projectId"])
          self.assertEqual(job.destination.dataset_id, table_ref["datasetId"])
          self.assertEqual(job.destination.table_id, table_ref["tableId"])
  
          sources = config.get("sourceTables")
          if sources is None:
              sources = [config["sourceTable"]]
          self.assertEqual(len(sources), len(job.sources))
          for table_ref, table in zip(sources, job.sources):
              self.assertEqual(table.project, table_ref["projectId"])
              self.assertEqual(table.dataset_id, table_ref["datasetId"])
              self.assertEqual(table.table_id, table_ref["tableId"])
  
          if "createDisposition" in config:
              self.assertEqual(job.create_disposition, config["createDisposition"])
          else:
              self.assertIsNone(job.create_disposition)
  
          if "writeDisposition" in config:
              self.assertEqual(job.write_disposition, config["writeDisposition"])
          else:
              self.assertIsNone(job.write_disposition)
  
          if "destinationEncryptionConfiguration" in config:
              self.assertIsNotNone(job.destination_encryption_configuration)
              self.assertEqual(
                  job.destination_encryption_configuration.kms_key_name,
                  config["destinationEncryptionConfiguration"]["kmsKeyName"],
              )
          else:
              self.assertIsNone(job.destination_encryption_configuration)
  
      def test_ctor(self):
          client = _make_client(project=self.PROJECT)
          source = self._table_ref(self.SOURCE_TABLE)
          destination = self._table_ref(self.DESTINATION_TABLE)
          job = self._make_one(self.JOB_ID, [source], destination, client)
          self.assertEqual(job.destination, destination)
          self.assertEqual(job.sources, [source])
          self.assertIs(job._client, client)
          self.assertEqual(job.job_type, self.JOB_TYPE)
          self.assertEqual(job.path, "/projects/%s/jobs/%s" % (self.PROJECT, self.JOB_ID))
  
          self._verifyInitialReadonlyProperties(job)
  
          # set/read from resource['configuration']['copy']
          self.assertIsNone(job.create_disposition)
          self.assertIsNone(job.write_disposition)
          self.assertIsNone(job.destination_encryption_configuration)
  
      def test_from_api_repr_missing_identity(self):
          self._setUpConstants()
          client = _make_client(project=self.PROJECT)
          RESOURCE = {}
          klass = self._get_target_class()
          with self.assertRaises(KeyError):
              klass.from_api_repr(RESOURCE, client=client)
  
      def test_from_api_repr_missing_config(self):
          self._setUpConstants()
          client = _make_client(project=self.PROJECT)
          RESOURCE = {
              "id": "%s:%s" % (self.PROJECT, self.DS_ID),
              "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
          }
          klass = self._get_target_class()
          with self.assertRaises(KeyError):
              klass.from_api_repr(RESOURCE, client=client)
  
      def test_from_api_repr_bare(self):
          self._setUpConstants()
          client = _make_client(project=self.PROJECT)
          RESOURCE = {
              "id": self.JOB_ID,
              "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
              "configuration": {
                  "copy": {
                      "sourceTables": [
                          {
                              "projectId": self.PROJECT,
                              "datasetId": self.DS_ID,
                              "tableId": self.SOURCE_TABLE,
                          }
                      ],
                      "destinationTable": {
                          "projectId": self.PROJECT,
                          "datasetId": self.DS_ID,
                          "tableId": self.DESTINATION_TABLE,
                      },
                  }
              },
          }
          klass = self._get_target_class()
          job = klass.from_api_repr(RESOURCE, client=client)
          self.assertIs(job._client, client)
          self._verifyResourceProperties(job, RESOURCE)
  
      def test_from_api_with_encryption(self):
          self._setUpConstants()
          client = _make_client(project=self.PROJECT)
          RESOURCE = {
              "id": self.JOB_ID,
              "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
              "configuration": {
                  "copy": {
                      "sourceTables": [
                          {
                              "projectId": self.PROJECT,
                              "datasetId": self.DS_ID,
                              "tableId": self.SOURCE_TABLE,
                          }
                      ],
                      "destinationTable": {
                          "projectId": self.PROJECT,
                          "datasetId": self.DS_ID,
                          "tableId": self.DESTINATION_TABLE,
                      },
                      "destinationEncryptionConfiguration": {
                          "kmsKeyName": self.KMS_KEY_NAME
                      },
                  }
              },
          }
          klass = self._get_target_class()
          job = klass.from_api_repr(RESOURCE, client=client)
          self.assertIs(job._client, client)
          self._verifyResourceProperties(job, RESOURCE)
  
      def test_from_api_repr_w_sourcetable(self):
          self._setUpConstants()
          client = _make_client(project=self.PROJECT)
          RESOURCE = {
              "id": self.JOB_ID,
              "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
              "configuration": {
                  "copy": {
                      "sourceTable": {
                          "projectId": self.PROJECT,
                          "datasetId": self.DS_ID,
                          "tableId": self.SOURCE_TABLE,
                      },
                      "destinationTable": {
                          "projectId": self.PROJECT,
                          "datasetId": self.DS_ID,
                          "tableId": self.DESTINATION_TABLE,
                      },
                  }
              },
          }
          klass = self._get_target_class()
          job = klass.from_api_repr(RESOURCE, client=client)
          self.assertIs(job._client, client)
          self._verifyResourceProperties(job, RESOURCE)
  
      def test_from_api_repr_wo_sources(self):
          self._setUpConstants()
          client = _make_client(project=self.PROJECT)
          RESOURCE = {
              "id": self.JOB_ID,
              "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
              "configuration": {
                  "copy": {
                      "destinationTable": {
                          "projectId": self.PROJECT,
                          "datasetId": self.DS_ID,
                          "tableId": self.DESTINATION_TABLE,
                      }
                  }
              },
          }
          klass = self._get_target_class()
          job = klass.from_api_repr(RESOURCE, client=client)
          with self.assertRaises(KeyError):
              _ = job.sources
  
      def test_from_api_repr_w_properties(self):
          from google.cloud.bigquery.job import CreateDisposition
  
          client = _make_client(project=self.PROJECT)
          RESOURCE = self._make_resource()
          copy_config = RESOURCE["configuration"]["copy"]
          copy_config["createDisposition"] = CreateDisposition.CREATE_IF_NEEDED
          klass = self._get_target_class()
          job = klass.from_api_repr(RESOURCE, client=client)
          self.assertIs(job._client, client)
          self._verifyResourceProperties(job, RESOURCE)
  
      def test_begin_w_bound_client(self):
          PATH = "/projects/%s/jobs" % (self.PROJECT,)
          RESOURCE = self._make_resource()
          # Ensure None for missing server-set props
          del RESOURCE["statistics"]["creationTime"]
          del RESOURCE["etag"]
          del RESOURCE["selfLink"]
          del RESOURCE["user_email"]
          conn = _make_connection(RESOURCE)
          client = _make_client(project=self.PROJECT, connection=conn)
          source = self._table_ref(self.SOURCE_TABLE)
          destination = self._table_ref(self.DESTINATION_TABLE)
          job = self._make_one(self.JOB_ID, [source], destination, client)
          with mock.patch(
              "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
          ) as final_attributes:
              job._begin()
  
          final_attributes.assert_called_with({"path": PATH}, client, job)
  
          conn.api_request.assert_called_once_with(
              method="POST",
              path=PATH,
              data={
                  "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
                  "configuration": {
                      "copy": {
                          "sourceTables": [
                              {
                                  "projectId": self.PROJECT,
                                  "datasetId": self.DS_ID,
                                  "tableId": self.SOURCE_TABLE,
                              }
                          ],
                          "destinationTable": {
                              "projectId": self.PROJECT,
                              "datasetId": self.DS_ID,
                              "tableId": self.DESTINATION_TABLE,
                          },
                      }
                  },
              },
              timeout=None,
          )
          self._verifyResourceProperties(job, RESOURCE)
  
      def test_begin_w_alternate_client(self):
          from google.cloud.bigquery.job import CopyJobConfig
  
          from google.cloud.bigquery.job import CreateDisposition
          from google.cloud.bigquery.job import WriteDisposition
  
          PATH = "/projects/%s/jobs" % (self.PROJECT,)
          RESOURCE = self._make_resource(ended=True)
          COPY_CONFIGURATION = {
              "sourceTables": [
                  {
                      "projectId": self.PROJECT,
                      "datasetId": self.DS_ID,
                      "tableId": self.SOURCE_TABLE,
                  }
              ],
              "destinationTable": {
                  "projectId": self.PROJECT,
                  "datasetId": self.DS_ID,
                  "tableId": self.DESTINATION_TABLE,
              },
              "createDisposition": CreateDisposition.CREATE_NEVER,
              "writeDisposition": WriteDisposition.WRITE_TRUNCATE,
          }
          RESOURCE["configuration"]["copy"] = COPY_CONFIGURATION
          conn1 = _make_connection()
          client1 = _make_client(project=self.PROJECT, connection=conn1)
          conn2 = _make_connection(RESOURCE)
          client2 = _make_client(project=self.PROJECT, connection=conn2)
          source = self._table_ref(self.SOURCE_TABLE)
          destination = self._table_ref(self.DESTINATION_TABLE)
          config = CopyJobConfig()
          config.create_disposition = CreateDisposition.CREATE_NEVER
          config.write_disposition = WriteDisposition.WRITE_TRUNCATE
          job = self._make_one(self.JOB_ID, [source], destination, client1, config)
          with mock.patch(
              "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
          ) as final_attributes:
              job._begin(client=client2)
  
          final_attributes.assert_called_with({"path": PATH}, client2, job)
  
          conn1.api_request.assert_not_called()
          conn2.api_request.assert_called_once_with(
              method="POST",
              path=PATH,
              data={
                  "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
                  "configuration": {"copy": COPY_CONFIGURATION},
              },
              timeout=None,
          )
          self._verifyResourceProperties(job, RESOURCE)
  
      def test_exists_miss_w_bound_client(self):
          PATH = "/projects/%s/jobs/%s" % (self.PROJECT, self.JOB_ID)
          conn = _make_connection()
          client = _make_client(project=self.PROJECT, connection=conn)
  
          source = self._table_ref(self.SOURCE_TABLE)
          destination = self._table_ref(self.DESTINATION_TABLE)
          job = self._make_one(self.JOB_ID, [source], destination, client)
          with mock.patch(
              "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
          ) as final_attributes:
              self.assertFalse(job.exists())
  
          final_attributes.assert_called_with({"path": PATH}, client, job)
  
          conn.api_request.assert_called_once_with(
              method="GET", path=PATH, query_params={"fields": "id"}, timeout=None
          )
  
      def test_exists_hit_w_alternate_client(self):
          PATH = "/projects/%s/jobs/%s" % (self.PROJECT, self.JOB_ID)
          conn1 = _make_connection()
          client1 = _make_client(project=self.PROJECT, connection=conn1)
          conn2 = _make_connection({})
          client2 = _make_client(project=self.PROJECT, connection=conn2)
          source = self._table_ref(self.SOURCE_TABLE)
          destination = self._table_ref(self.DESTINATION_TABLE)
          job = self._make_one(self.JOB_ID, [source], destination, client1)
          with mock.patch(
              "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
          ) as final_attributes:
              self.assertTrue(job.exists(client=client2))
  
          final_attributes.assert_called_with({"path": PATH}, client2, job)
  
          conn1.api_request.assert_not_called()
          conn2.api_request.assert_called_once_with(
              method="GET", path=PATH, query_params={"fields": "id"}, timeout=None
          )
  
      def test_reload_w_bound_client(self):
          PATH = "/projects/%s/jobs/%s" % (self.PROJECT, self.JOB_ID)
          RESOURCE = self._make_resource()
          conn = _make_connection(RESOURCE)
          client = _make_client(project=self.PROJECT, connection=conn)
          source = self._table_ref(self.SOURCE_TABLE)
          destination = self._table_ref(self.DESTINATION_TABLE)
          job = self._make_one(self.JOB_ID, [source], destination, client)
          with mock.patch(
              "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
          ) as final_attributes:
              job.reload()
  
          final_attributes.assert_called_with({"path": PATH}, client, job)
  
          conn.api_request.assert_called_once_with(
              method="GET", path=PATH, query_params={}, timeout=None
          )
          self._verifyResourceProperties(job, RESOURCE)
  
      def test_reload_w_alternate_client(self):
          PATH = "/projects/%s/jobs/%s" % (self.PROJECT, self.JOB_ID)
          RESOURCE = self._make_resource()
          conn1 = _make_connection()
          client1 = _make_client(project=self.PROJECT, connection=conn1)
          conn2 = _make_connection(RESOURCE)
          client2 = _make_client(project=self.PROJECT, connection=conn2)
          source = self._table_ref(self.SOURCE_TABLE)
          destination = self._table_ref(self.DESTINATION_TABLE)
          job = self._make_one(self.JOB_ID, [source], destination, client1)
          with mock.patch(
              "google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"
          ) as final_attributes:
              job.reload(client=client2)
  
          final_attributes.assert_called_with({"path": PATH}, client2, job)
  
          conn1.api_request.assert_not_called()
          conn2.api_request.assert_called_once_with(
              method="GET", path=PATH, query_params={}, timeout=None
          )
          self._verifyResourceProperties(job, RESOURCE)
  