sdks/python/apache_beam/io/gcp/bigquery_biglake_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

"""Unit tests for BigQuery BigLake configuration."""

import unittest
from unittest import mock

from apache_beam.io.gcp import bigquery


@mock.patch('apache_beam.io.gcp.bigquery.BeamJarExpansionService')
class BigQueryBigLakeTest(unittest.TestCase):
  """Test BigLake configuration support in BigQuery Storage Write API."""
  def test_storage_write_to_bigquery_with_biglake_config(
      self, mock_expansion_service):
    """Test that StorageWriteToBigQuery accepts bigLakeConfiguration."""
    big_lake_config = {
        'connectionId': (
            'projects/test-project/locations/us/connections/test-connection'),
        'storageUri': 'gs://test-bucket/test-path',
        'fileFormat': 'parquet',
        'tableFormat': 'iceberg'
    }

    # Test that the constructor accepts the bigLakeConfiguration parameter
    transform = bigquery.StorageWriteToBigQuery(
        table='test-project:test_dataset.test_table',
        big_lake_configuration=big_lake_config)

    # Verify the configuration is stored
    self.assertEqual(transform._big_lake_configuration, big_lake_config)

  def test_storage_write_to_bigquery_without_biglake_config(
      self, mock_expansion_service):
    """Test that StorageWriteToBigQuery works without bigLakeConfiguration."""
    transform = bigquery.StorageWriteToBigQuery(
        table='test-project:test_dataset.test_table')

    # Verify the configuration is None by default
    self.assertIsNone(transform._big_lake_configuration)

  def test_biglake_config_passed_to_external_transform(
      self, mock_expansion_service):
    """Test that StorageWriteToBigQuery accepts bigLakeConfiguration."""
    big_lake_config = {
        'connection_id': 'projects/my-project/locations/us/connections/my-conn',
        'table_format': 'ICEBERG'
    }

    # Mock the expansion service to avoid JAR dependency
    mock_expansion_service.return_value = mock.MagicMock()

    # Create the transform
    transform = bigquery.StorageWriteToBigQuery(
        table='my-project:my_dataset.my_table',
        big_lake_configuration=big_lake_config)

    # Verify the big_lake_configuration is stored correctly
    self.assertEqual(transform._big_lake_configuration, big_lake_config)

    # Verify that the transform has the expected identifier
    self.assertEqual(
        transform.IDENTIFIER,
        "beam:schematransform:org.apache.beam:bigquery_storage_write:v2")

    # Verify that the expansion service was created (mocked)
    mock_expansion_service.assert_called_once_with(
        'sdks:java:io:google-cloud-platform:expansion-service:build')

  def test_biglake_config_validation(self, mock_expansion_service):
    """Test validation of bigLakeConfiguration parameters."""
    # Test with minimal required configuration
    minimal_config = {
        'connectionId': (
            'projects/test-project/locations/us/connections/test-connection'),
        'storageUri': 'gs://test-bucket/test-path'
    }

    transform = bigquery.StorageWriteToBigQuery(
        table='test-project:test_dataset.test_table',
        big_lake_configuration=minimal_config)

    self.assertEqual(transform._big_lake_configuration, minimal_config)

    # Test with full configuration
    full_config = {
        'connectionId': (
            'projects/test-project/locations/us/connections/test-connection'),
        'storageUri': 'gs://test-bucket/test-path',
        'fileFormat': 'parquet',
        'tableFormat': 'iceberg'
    }

    transform = bigquery.StorageWriteToBigQuery(
        table='test-project:test_dataset.test_table',
        big_lake_configuration=full_config)

    self.assertEqual(transform._big_lake_configuration, full_config)


if __name__ == '__main__':
  unittest.main()