content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  
  import unittest
  from unittest.mock import MagicMock, patch
  
  from pyhocon import ConfigFactory
  
  from databuilder.models.table_metadata import ColumnMetadata, TableMetadata
  from databuilder.models.type_metadata import (
      ArrayTypeMetadata, ScalarTypeMetadata, TypeMetadata,
  )
  from databuilder.transformer.complex_type_transformer import PARSING_FUNCTION, ComplexTypeTransformer
  
  
  class TestComplexTypeTransformer(unittest.TestCase):
      def test_invalid_parsing_function_missing_module(self) -> None:
          transformer = ComplexTypeTransformer()
          config = ConfigFactory.from_dict({
              PARSING_FUNCTION: 'invalid_function',
          })
          with self.assertRaises(Exception):
              transformer.init(conf=config)
  
      def test_invalid_parsing_function_invalid_module(self) -> None:
          transformer = ComplexTypeTransformer()
          config = ConfigFactory.from_dict({
              PARSING_FUNCTION: 'invalid_module.invalid_function',
          })
          with self.assertRaises(Exception):
              transformer.init(conf=config)
  
      def test_invalid_parsing_function_invalid_function(self) -> None:
          transformer = ComplexTypeTransformer()
          config = ConfigFactory.from_dict({
              PARSING_FUNCTION: 'databuilder.utils.hive_complex_type_parser.invalid_function',
          })
          with self.assertRaises(Exception):
              transformer.init(conf=config)
  
      def test_hive_parser_with_failures(self) -> None:
          transformer = ComplexTypeTransformer()
          config = ConfigFactory.from_dict({
              PARSING_FUNCTION: 'databuilder.utils.hive_complex_type_parser.parse_hive_type',
          })
          transformer.init(conf=config)
  
          column = ColumnMetadata('col1', 'array type', 'array<array<int>>', 0)
          table_metadata = TableMetadata(
              'hive',
              'gold',
              'test_schema',
              'test_table',
              'test_table',
              [column]
          )
  
          default_scalar_type = ScalarTypeMetadata(name='col1',
                                                   parent=column,
                                                   type_str='array<array<int>>')
  
          with patch.object(transformer, '_parsing_function') as mock:
              mock.side_effect = MagicMock(side_effect=Exception('Could not parse'))
  
              result = transformer.transform(table_metadata)
  
              self.assertEqual(transformer.success_count, 0)
              self.assertEqual(transformer.failure_count, 1)
              for actual in result.columns:
                  self.assertEqual(actual.get_type_metadata(), default_scalar_type)
  
      def test_hive_parser_usage(self) -> None:
          transformer = ComplexTypeTransformer()
          config = ConfigFactory.from_dict({
              PARSING_FUNCTION: 'databuilder.utils.hive_complex_type_parser.parse_hive_type',
          })
          transformer.init(conf=config)
  
          column = ColumnMetadata('col1', 'array type', 'array<array<int>>', 0)
          table_metadata = TableMetadata(
              'hive',
              'gold',
              'test_schema',
              'test_table',
              'test_table',
              [column]
          )
          array_type = ArrayTypeMetadata(name='col1',
                                         parent=column,
                                         type_str='array<array<int>>')
          inner_array = ArrayTypeMetadata(name='_inner_',
                                          parent=array_type,
                                          type_str='array<int>')
  
          array_type.array_inner_type = inner_array
  
          result = transformer.transform(table_metadata)
  
          for actual in result.columns:
              self.assertTrue(isinstance(actual.get_type_metadata(), TypeMetadata))
              self.assertEqual(actual.get_type_metadata(), array_type)
              self.assertEqual(transformer.success_count, 1)
              self.assertEqual(transformer.failure_count, 0)
  
      def test_trino_parser_usage(self) -> None:
          transformer = ComplexTypeTransformer()
          config = ConfigFactory.from_dict({
              PARSING_FUNCTION: 'databuilder.utils.trino_complex_type_parser.parse_trino_type',
          })
          transformer.init(conf=config)
  
          column = ColumnMetadata('col1', 'array type', 'array(array(int))', 0)
          table_metadata = TableMetadata(
              'trino',
              'gold',
              'test_schema',
              'test_table',
              'test_table',
              [column]
          )
          array_type = ArrayTypeMetadata(name='col1',
                                         parent=column,
                                         type_str='array(array(int))')
          inner_array = ArrayTypeMetadata(name='_inner_',
                                          parent=array_type,
                                          type_str='array(int)')
  
          array_type.array_inner_type = inner_array
  
          result = transformer.transform(table_metadata)
  
          for actual in result.columns:
              self.assertTrue(isinstance(actual.get_type_metadata(), TypeMetadata))
              self.assertEqual(actual.get_type_metadata(), array_type)
              self.assertEqual(transformer.success_count, 1)
              self.assertEqual(transformer.failure_count, 0)
  
  
  if __name__ == '__main__':
      unittest.main()
  