content: # Copyright Contributors to the Amundsen project.
  # SPDX-License-Identifier: Apache-2.0
  import unittest
  from typing import Any
  
  from mock import MagicMock, Mock
  from pyhocon import ConfigFactory
  
  from databuilder import Scoped
  from databuilder.extractor.dashboard.apache_superset.apache_superset_table_extractor import ApacheSupersetTableExtractor
  from databuilder.models.dashboard.dashboard_table import DashboardTable
  
  dataset_data_response_1 = {
      'result': {
          'sql': None,
          'table_name': 'table_name',
          'database': {
              'id': 1
          }
      }
  }
  
  dataset_objects_data_response_1 = {
      'dashboards': {
          'result': [
              {
                  'id': 2
              }
          ]
      }
  }
  
  database_data_response_1 = {
      'result': {
          'sqlalchemy_uri': 'postgresql://localhost:5432/db_name'
      }
  }
  
  dataset_data_response_2 = {
      'result': {
          'sql': None,
          'table_name': 'table_name_2',
          'database': {
              'id': 3
          }
      }
  }
  
  dataset_objects_data_response_2 = {
      'dashboards': {
          'result': [
              {
                  'id': 2
              }
          ]
      }
  }
  
  database_data_response_2 = {
      'result': {
          'sqlalchemy_uri': 'postgresql://localhost:5432/db_name_2'
      }
  }
  
  
  class TestApacheSupersetTableExtractor(unittest.TestCase):
      def setUp(self) -> None:
          config = ConfigFactory.from_dict({
              'extractor.apache_superset.dashboard_group_id': '1',
              'extractor.apache_superset.dashboard_group_name': 'dashboard group',
              'extractor.apache_superset.dashboard_group_description': 'dashboard group description',
              'extractor.apache_superset.cluster': 'gold',
              'extractor.apache_superset.apache_superset_security_settings_dict': dict(username='admin',
                                                                                       password='admin',
                                                                                       provider='db')
          })
  
          self.config = config
  
      def _get_extractor(self) -> Any:
          extractor = self._extractor_class()
          extractor.authenticate = MagicMock()
          extractor.init(Scoped.get_scoped_conf(conf=self.config, scope=extractor.get_scope()))
  
          return extractor
  
      @property
      def _extractor_class(self) -> Any:
          return ApacheSupersetTableExtractor
  
      def test_extractor(self) -> None:
          extractor = self._get_extractor()
  
          extractor.execute_query = Mock(side_effect=[{'ids': [2, 3]}, {'ids': []},
                                                      dataset_data_response_1, dataset_objects_data_response_1,
                                                      dataset_data_response_2, dataset_objects_data_response_2,
                                                      database_data_response_1, database_data_response_2])
  
          record = extractor.extract()
  
          self.assertIsInstance(record, DashboardTable)
          self.assertEquals(record._dashboard_id, '2')
          self.assertSetEqual(record._table_ids,
                              {'postgres://gold.db_name/table_name', 'postgres://gold.db_name_2/table_name_2'})
  