Unlock 30% off on Manual Testing Annual Plans this Holiday Season.

Copied to Clipboard!

How to use get_client method in locust

Best Python code snippet using locust

test_cli.py

Source:test_cli.py

...13@pytest.fixture14def runner():15    return CliRunner()16@patch.object(EcsClient, '__init__')17def test_get_client(ecs_client):18    ecs_client.return_value = None19    client = get_client('access_key_id', 'secret_access_key', 'region', 'profile')20    ecs_client.assert_called_once_with('access_key_id', 'secret_access_key', 'region', 'profile')21    assert isinstance(client, EcsClient)22def test_ecs(runner):23    result = runner.invoke(cli.ecs)24    assert result.exit_code == 025    assert not result.exception26    assert 'Usage: ecs [OPTIONS] COMMAND [ARGS]' in result.output27    assert '  deploy  ' in result.output28    assert '  scale   ' in result.output29@patch('ecs_deploy.cli.get_client')30def test_deploy_without_credentials(get_client, runner):31    get_client.return_value = EcsTestClient()32    result = runner.invoke(cli.deploy, (CLUSTER_NAME, SERVICE_NAME))33    assert result.exit_code == 1...

test_gcp_spanner_hook.py

Source:test_gcp_spanner_hook.py

1# -*- coding: utf-8 -*-2#3# Licensed to the Apache Software Foundation (ASF) under one4# or more contributor license agreements.  See the NOTICE file5# distributed with this work for additional information6# regarding copyright ownership.  The ASF licenses this file7# to you under the Apache License, Version 2.0 (the8# "License"); you may not use this file except in compliance9# with the License.  You may obtain a copy of the License at10#11#   http://www.apache.org/licenses/LICENSE-2.012#13# Unless required by applicable law or agreed to in writing,14# software distributed under the License is distributed on an15# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY16# KIND, either express or implied.  See the License for the17# specific language governing permissions and limitations18# under the License.19import unittest20from airflow import AirflowException21from airflow.contrib.hooks.gcp_spanner_hook import CloudSpannerHook22from tests.contrib.utils.base_gcp_mock import mock_base_gcp_hook_no_default_project_id, \23    GCP_PROJECT_ID_HOOK_UNIT_TEST, mock_base_gcp_hook_default_project_id24from tests.compat import mock25SPANNER_INSTANCE = 'instance'26SPANNER_CONFIGURATION = 'configuration'27SPANNER_DATABASE = 'database-name'28class TestGcpSpannerHookDefaultProjectId(unittest.TestCase):29    def setUp(self):30        with mock.patch('airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__',31                        new=mock_base_gcp_hook_default_project_id):32            self.spanner_hook_default_project_id = CloudSpannerHook(gcp_conn_id='test')33    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')34    def test_get_existing_instance(self, get_client):35        instance_method = get_client.return_value.instance36        instance_exists_method = instance_method.return_value.exists37        instance_exists_method.return_value = True38        res = self.spanner_hook_default_project_id.get_instance(instance_id=SPANNER_INSTANCE,39                                                                project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST)40        get_client.assert_called_once_with(project_id='example-project')41        instance_method.assert_called_once_with(instance_id='instance')42        self.assertIsNotNone(res)43    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')44    def test_get_existing_instance_overridden_project_id(self, get_client):45        instance_method = get_client.return_value.instance46        instance_exists_method = instance_method.return_value.exists47        instance_exists_method.return_value = True48        res = self.spanner_hook_default_project_id.get_instance(instance_id=SPANNER_INSTANCE,49                                                                project_id='new-project')50        get_client.assert_called_once_with(project_id='new-project')51        instance_method.assert_called_once_with(instance_id='instance')52        self.assertIsNotNone(res)53    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')54    def test_create_instance(self, get_client):55        instance_method = get_client.return_value.instance56        create_method = instance_method.return_value.create57        create_method.return_value = False58        res = self.spanner_hook_default_project_id.create_instance(59            instance_id=SPANNER_INSTANCE,60            configuration_name=SPANNER_CONFIGURATION,61            node_count=1,62            display_name=SPANNER_DATABASE)63        get_client.assert_called_once_with(project_id='example-project')64        instance_method.assert_called_once_with(65            instance_id='instance',66            configuration_name='configuration',67            display_name='database-name',68            node_count=1)69        self.assertIsNone(res)70    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')71    def test_create_instance_overridden_project_id(self, get_client):72        instance_method = get_client.return_value.instance73        create_method = instance_method.return_value.create74        create_method.return_value = False75        res = self.spanner_hook_default_project_id.create_instance(76            project_id='new-project',77            instance_id=SPANNER_INSTANCE,78            configuration_name=SPANNER_CONFIGURATION,79            node_count=1,80            display_name=SPANNER_DATABASE)81        get_client.assert_called_once_with(project_id='new-project')82        instance_method.assert_called_once_with(83            instance_id='instance',84            configuration_name='configuration',85            display_name='database-name',86            node_count=1)87        self.assertIsNone(res)88    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')89    def test_update_instance(self, get_client):90        instance_method = get_client.return_value.instance91        instance_exists_method = instance_method.return_value.exists92        instance_exists_method.return_value = True93        update_method = instance_method.return_value.update94        update_method.return_value = False95        res = self.spanner_hook_default_project_id.update_instance(96            instance_id=SPANNER_INSTANCE,97            configuration_name=SPANNER_CONFIGURATION,98            node_count=2,99            display_name=SPANNER_DATABASE)100        get_client.assert_called_once_with(project_id='example-project')101        instance_method.assert_called_once_with(102            instance_id='instance', configuration_name='configuration', display_name='database-name',103            node_count=2)104        update_method.assert_called_with()105        self.assertIsNone(res)106    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')107    def test_update_instance_overridden_project_id(self, get_client):108        instance_method = get_client.return_value.instance109        instance_exists_method = instance_method.return_value.exists110        instance_exists_method.return_value = True111        update_method = instance_method.return_value.update112        update_method.return_value = False113        res = self.spanner_hook_default_project_id.update_instance(114            project_id='new-project',115            instance_id=SPANNER_INSTANCE,116            configuration_name=SPANNER_CONFIGURATION,117            node_count=2,118            display_name=SPANNER_DATABASE)119        get_client.assert_called_once_with(project_id='new-project')120        instance_method.assert_called_once_with(121            instance_id='instance', configuration_name='configuration', display_name='database-name',122            node_count=2)123        update_method.assert_called_with()124        self.assertIsNone(res)125    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')126    def test_delete_instance(self, get_client):127        instance_method = get_client.return_value.instance128        instance_exists_method = instance_method.return_value.exists129        instance_exists_method.return_value = True130        delete_method = instance_method.return_value.delete131        delete_method.return_value = False132        res = self.spanner_hook_default_project_id.delete_instance(133            instance_id=SPANNER_INSTANCE)134        get_client.assert_called_once_with(project_id='example-project')135        instance_method.assert_called_once_with(136            'instance')137        delete_method.assert_called_with()138        self.assertIsNone(res)139    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')140    def test_delete_instance_overridden_project_id(self, get_client):141        instance_method = get_client.return_value.instance142        instance_exists_method = instance_method.return_value.exists143        instance_exists_method.return_value = True144        delete_method = instance_method.return_value.delete145        delete_method.return_value = False146        res = self.spanner_hook_default_project_id.delete_instance(147            project_id='new-project',148            instance_id=SPANNER_INSTANCE)149        get_client.assert_called_once_with(project_id='new-project')150        instance_method.assert_called_once_with(151            'instance')152        delete_method.assert_called_with()153        self.assertIsNone(res)154    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')155    def test_get_database(self, get_client):156        instance_method = get_client.return_value.instance157        instance_exists_method = instance_method.return_value.exists158        instance_exists_method.return_value = True159        database_method = instance_method.return_value.database160        database_exists_method = instance_method.return_value.exists161        database_exists_method.return_value = True162        res = self.spanner_hook_default_project_id.get_database(163            instance_id=SPANNER_INSTANCE,164            database_id=SPANNER_DATABASE)165        get_client.assert_called_once_with(project_id='example-project')166        instance_method.assert_called_once_with(instance_id='instance')167        database_method.assert_called_with(database_id='database-name')168        database_exists_method.assert_called_with()169        self.assertIsNotNone(res)170    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')171    def test_get_database_overridden_project_id(self, get_client):172        instance_method = get_client.return_value.instance173        instance_exists_method = instance_method.return_value.exists174        instance_exists_method.return_value = True175        database_method = instance_method.return_value.database176        database_exists_method = instance_method.return_value.exists177        database_exists_method.return_value = True178        res = self.spanner_hook_default_project_id.get_database(179            project_id='new-project',180            instance_id=SPANNER_INSTANCE,181            database_id=SPANNER_DATABASE)182        get_client.assert_called_once_with(project_id='new-project')183        instance_method.assert_called_once_with(instance_id='instance')184        database_method.assert_called_with(database_id='database-name')185        database_exists_method.assert_called_with()186        self.assertIsNotNone(res)187    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')188    def test_create_database(self, get_client):189        instance_method = get_client.return_value.instance190        instance_exists_method = instance_method.return_value.exists191        instance_exists_method.return_value = True192        database_method = instance_method.return_value.database193        database_create_method = database_method.return_value.create194        res = self.spanner_hook_default_project_id.create_database(195            instance_id=SPANNER_INSTANCE,196            database_id=SPANNER_DATABASE,197            ddl_statements=[])198        get_client.assert_called_once_with(project_id='example-project')199        instance_method.assert_called_once_with(instance_id='instance')200        database_method.assert_called_with(database_id='database-name', ddl_statements=[])201        database_create_method.assert_called_with()202        self.assertIsNone(res)203    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')204    def test_create_database_overridden_project_id(self, get_client):205        instance_method = get_client.return_value.instance206        instance_exists_method = instance_method.return_value.exists207        instance_exists_method.return_value = True208        database_method = instance_method.return_value.database209        database_create_method = database_method.return_value.create210        res = self.spanner_hook_default_project_id.create_database(211            project_id='new-project',212            instance_id=SPANNER_INSTANCE,213            database_id=SPANNER_DATABASE,214            ddl_statements=[])215        get_client.assert_called_once_with(project_id='new-project')216        instance_method.assert_called_once_with(instance_id='instance')217        database_method.assert_called_with(database_id='database-name', ddl_statements=[])218        database_create_method.assert_called_with()219        self.assertIsNone(res)220    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')221    def test_update_database(self, get_client):222        instance_method = get_client.return_value.instance223        instance_exists_method = instance_method.return_value.exists224        instance_exists_method.return_value = True225        database_method = instance_method.return_value.database226        database_update_ddl_method = database_method.return_value.update_ddl227        res = self.spanner_hook_default_project_id.update_database(228            instance_id=SPANNER_INSTANCE,229            database_id=SPANNER_DATABASE,230            ddl_statements=[])231        get_client.assert_called_once_with(project_id='example-project')232        instance_method.assert_called_once_with(instance_id='instance')233        database_method.assert_called_with(database_id='database-name')234        database_update_ddl_method.assert_called_with(ddl_statements=[], operation_id=None)235        self.assertIsNone(res)236    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')237    def test_update_database_overridden_project_id(self, get_client):238        instance_method = get_client.return_value.instance239        instance_exists_method = instance_method.return_value.exists240        instance_exists_method.return_value = True241        database_method = instance_method.return_value.database242        database_update_ddl_method = database_method.return_value.update_ddl243        res = self.spanner_hook_default_project_id.update_database(244            project_id='new-project',245            instance_id=SPANNER_INSTANCE,246            database_id=SPANNER_DATABASE,247            ddl_statements=[])248        get_client.assert_called_once_with(project_id='new-project')249        instance_method.assert_called_once_with(instance_id='instance')250        database_method.assert_called_with(database_id='database-name')251        database_update_ddl_method.assert_called_with(ddl_statements=[], operation_id=None)252        self.assertIsNone(res)253    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')254    def test_delete_database(self, get_client):255        instance_method = get_client.return_value.instance256        instance_exists_method = instance_method.return_value.exists257        instance_exists_method.return_value = True258        database_method = instance_method.return_value.database259        database_drop_method = database_method.return_value.drop260        database_exists_method = database_method.return_value.exists261        database_exists_method.return_value = True262        res = self.spanner_hook_default_project_id.delete_database(263            instance_id=SPANNER_INSTANCE,264            database_id=SPANNER_DATABASE)265        get_client.assert_called_once_with(project_id='example-project')266        instance_method.assert_called_once_with(instance_id='instance')267        database_method.assert_called_with(database_id='database-name')268        database_exists_method.assert_called_with()269        database_drop_method.assert_called_with()270        self.assertIsNone(res)271    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')272    def test_delete_database_overridden_project_id(self, get_client):273        instance_method = get_client.return_value.instance274        instance_exists_method = instance_method.return_value.exists275        instance_exists_method.return_value = True276        database_method = instance_method.return_value.database277        database_drop_method = database_method.return_value.drop278        database_exists_method = database_method.return_value.exists279        database_exists_method.return_value = True280        res = self.spanner_hook_default_project_id.delete_database(281            project_id='new-project',282            instance_id=SPANNER_INSTANCE,283            database_id=SPANNER_DATABASE)284        get_client.assert_called_once_with(project_id='new-project')285        instance_method.assert_called_once_with(instance_id='instance')286        database_method.assert_called_with(database_id='database-name')287        database_exists_method.assert_called_with()288        database_drop_method.assert_called_with()289        self.assertIsNone(res)290    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')291    def test_execute_dml(self, get_client):292        instance_method = get_client.return_value.instance293        instance_exists_method = instance_method.return_value.exists294        instance_exists_method.return_value = True295        database_method = instance_method.return_value.database296        run_in_transaction_method = database_method.return_value.run_in_transaction297        res = self.spanner_hook_default_project_id.execute_dml(298            instance_id=SPANNER_INSTANCE,299            database_id=SPANNER_DATABASE,300            queries='')301        get_client.assert_called_once_with(project_id='example-project')302        instance_method.assert_called_once_with(instance_id='instance')303        database_method.assert_called_with(database_id='database-name')304        run_in_transaction_method.assert_called_with(mock.ANY)305        self.assertIsNone(res)306    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')307    def test_execute_dml_overridden_project_id(self, get_client):308        instance_method = get_client.return_value.instance309        instance_exists_method = instance_method.return_value.exists310        instance_exists_method.return_value = True311        database_method = instance_method.return_value.database312        run_in_transaction_method = database_method.return_value.run_in_transaction313        res = self.spanner_hook_default_project_id.execute_dml(314            project_id='new-project',315            instance_id=SPANNER_INSTANCE,316            database_id=SPANNER_DATABASE,317            queries='')318        get_client.assert_called_once_with(project_id='new-project')319        instance_method.assert_called_once_with(instance_id='instance')320        database_method.assert_called_with(database_id='database-name')321        run_in_transaction_method.assert_called_with(mock.ANY)322        self.assertIsNone(res)323class TestGcpSpannerHookNoDefaultProjectID(unittest.TestCase):324    def setUp(self):325        with mock.patch('airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__',326                        new=mock_base_gcp_hook_no_default_project_id):327            self.spanner_hook_no_default_project_id = CloudSpannerHook(gcp_conn_id='test')328    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')329    def test_get_existing_instance_missing_project_id(self, get_client):330        instance_method = get_client.return_value.instance331        instance_exists_method = instance_method.return_value.exists332        instance_exists_method.return_value = True333        with self.assertRaises(AirflowException) as cm:334            self.spanner_hook_no_default_project_id.get_instance(instance_id=SPANNER_INSTANCE)335        get_client.assert_not_called()336        instance_method.assert_not_called()337        instance_exists_method.assert_not_called()338        err = cm.exception339        self.assertIn("The project id must be passed", str(err))340    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')341    def test_get_existing_instance_overridden_project_id(self, get_client):342        instance_method = get_client.return_value.instance343        instance_exists_method = instance_method.return_value.exists344        instance_exists_method.return_value = True345        res = self.spanner_hook_no_default_project_id.get_instance(instance_id=SPANNER_INSTANCE,346                                                                   project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST)347        get_client.assert_called_once_with(project_id='example-project')348        instance_method.assert_called_once_with(instance_id='instance')349        self.assertIsNotNone(res)350    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')351    def test_get_non_existing_instance(self, get_client):352        instance_method = get_client.return_value.instance353        instance_exists_method = instance_method.return_value.exists354        instance_exists_method.return_value = False355        res = self.spanner_hook_no_default_project_id.get_instance(instance_id=SPANNER_INSTANCE,356                                                                   project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST)357        get_client.assert_called_once_with(project_id='example-project')358        instance_method.assert_called_once_with(instance_id='instance')359        self.assertIsNone(res)360    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')361    def test_create_instance_missing_project_id(self, get_client):362        instance_method = get_client.return_value.instance363        create_method = instance_method.return_value.create364        create_method.return_value = False365        with self.assertRaises(AirflowException) as cm:366            self.spanner_hook_no_default_project_id.create_instance(367                instance_id=SPANNER_INSTANCE,368                configuration_name=SPANNER_CONFIGURATION,369                node_count=1,370                display_name=SPANNER_DATABASE)371            get_client.assert_called_once_with(project_id='example-project')372        get_client.assert_not_called()373        instance_method.assert_not_called()374        create_method.assert_not_called()375        err = cm.exception376        self.assertIn("The project id must be passed", str(err))377    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')378    def test_create_instance_overridden_project_id(self, get_client):379        instance_method = get_client.return_value.instance380        create_method = instance_method.return_value.create381        create_method.return_value = False382        res = self.spanner_hook_no_default_project_id.create_instance(383            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,384            instance_id=SPANNER_INSTANCE,385            configuration_name=SPANNER_CONFIGURATION,386            node_count=1,387            display_name=SPANNER_DATABASE)388        get_client.assert_called_once_with(project_id='example-project')389        instance_method.assert_called_once_with(390            instance_id='instance',391            configuration_name='configuration',392            display_name='database-name',393            node_count=1)394        self.assertIsNone(res)395    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')396    def test_update_instance_missing_project_id(self, get_client):397        instance_method = get_client.return_value.instance398        instance_exists_method = instance_method.return_value.exists399        instance_exists_method.return_value = True400        update_method = instance_method.return_value.update401        update_method.return_value = False402        with self.assertRaises(AirflowException) as cm:403            self.spanner_hook_no_default_project_id.update_instance(404                instance_id=SPANNER_INSTANCE,405                configuration_name=SPANNER_CONFIGURATION,406                node_count=2,407                display_name=SPANNER_DATABASE)408        get_client.assert_not_called()409        instance_method.assert_not_called()410        update_method.assert_not_called()411        err = cm.exception412        self.assertIn("The project id must be passed", str(err))413    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')414    def test_update_instance_overridden_project_id(self, get_client):415        instance_method = get_client.return_value.instance416        instance_exists_method = instance_method.return_value.exists417        instance_exists_method.return_value = True418        update_method = instance_method.return_value.update419        update_method.return_value = False420        res = self.spanner_hook_no_default_project_id.update_instance(421            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,422            instance_id=SPANNER_INSTANCE,423            configuration_name=SPANNER_CONFIGURATION,424            node_count=2,425            display_name=SPANNER_DATABASE)426        get_client.assert_called_once_with(project_id='example-project')427        instance_method.assert_called_once_with(428            instance_id='instance', configuration_name='configuration', display_name='database-name',429            node_count=2)430        update_method.assert_called_with()431        self.assertIsNone(res)432    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')433    def test_delete_instance_missing_project_id(self, get_client):434        instance_method = get_client.return_value.instance435        instance_exists_method = instance_method.return_value.exists436        instance_exists_method.return_value = True437        delete_method = instance_method.return_value.delete438        delete_method.return_value = False439        with self.assertRaises(AirflowException) as cm:440            self.spanner_hook_no_default_project_id.delete_instance(441                instance_id=SPANNER_INSTANCE)442        get_client.assert_not_called()443        instance_method.assert_not_called()444        delete_method.assert_not_called()445        err = cm.exception446        self.assertIn("The project id must be passed", str(err))447    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')448    def test_delete_instance_overridden_project_id(self, get_client):449        instance_method = get_client.return_value.instance450        instance_exists_method = instance_method.return_value.exists451        instance_exists_method.return_value = True452        delete_method = instance_method.return_value.delete453        delete_method.return_value = False454        res = self.spanner_hook_no_default_project_id.delete_instance(455            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,456            instance_id=SPANNER_INSTANCE)457        get_client.assert_called_once_with(project_id='example-project')458        instance_method.assert_called_once_with(459            'instance')460        delete_method.assert_called_with()461        self.assertIsNone(res)462    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')463    def test_get_database_missing_project_id(self, get_client):464        instance_method = get_client.return_value.instance465        instance_exists_method = instance_method.return_value.exists466        instance_exists_method.return_value = True467        database_method = instance_method.return_value.database468        database_exists_method = instance_method.return_value.exists469        database_exists_method.return_value = True470        with self.assertRaises(AirflowException) as cm:471            self.spanner_hook_no_default_project_id.get_database(472                instance_id=SPANNER_INSTANCE,473                database_id=SPANNER_DATABASE)474        get_client.assert_not_called()475        instance_method.assert_not_called()476        database_method.assert_not_called()477        err = cm.exception478        self.assertIn("The project id must be passed", str(err))479    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')480    def test_get_database_overridden_project_id(self, get_client):481        instance_method = get_client.return_value.instance482        instance_exists_method = instance_method.return_value.exists483        instance_exists_method.return_value = True484        database_method = instance_method.return_value.database485        database_exists_method = instance_method.return_value.exists486        database_exists_method.return_value = True487        res = self.spanner_hook_no_default_project_id.get_database(488            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,489            instance_id=SPANNER_INSTANCE,490            database_id=SPANNER_DATABASE)491        get_client.assert_called_once_with(project_id='example-project')492        instance_method.assert_called_once_with(instance_id='instance')493        database_method.assert_called_with(database_id='database-name')494        database_exists_method.assert_called_with()495        self.assertIsNotNone(res)496    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')497    def test_create_database_missing_project_id(self, get_client):498        instance_method = get_client.return_value.instance499        instance_exists_method = instance_method.return_value.exists500        instance_exists_method.return_value = True501        database_method = instance_method.return_value.database502        database_create_method = database_method.return_value.create503        with self.assertRaises(AirflowException) as cm:504            self.spanner_hook_no_default_project_id.create_database(505                instance_id=SPANNER_INSTANCE,506                database_id=SPANNER_DATABASE,507                ddl_statements=[])508        get_client.assert_not_called()509        instance_method.assert_not_called()510        database_method.assert_not_called()511        database_create_method.assert_not_called()512        err = cm.exception513        self.assertIn("The project id must be passed", str(err))514    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')515    def test_create_database_overridden_project_id(self, get_client):516        instance_method = get_client.return_value.instance517        instance_exists_method = instance_method.return_value.exists518        instance_exists_method.return_value = True519        database_method = instance_method.return_value.database520        database_create_method = database_method.return_value.create521        res = self.spanner_hook_no_default_project_id.create_database(522            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,523            instance_id=SPANNER_INSTANCE,524            database_id=SPANNER_DATABASE,525            ddl_statements=[])526        get_client.assert_called_once_with(project_id='example-project')527        instance_method.assert_called_once_with(instance_id='instance')528        database_method.assert_called_with(database_id='database-name', ddl_statements=[])529        database_create_method.assert_called_with()530        self.assertIsNone(res)531    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')532    def test_update_database_missing_project_id(self, get_client):533        instance_method = get_client.return_value.instance534        instance_exists_method = instance_method.return_value.exists535        instance_exists_method.return_value = True536        database_method = instance_method.return_value.database537        database_update_method = database_method.return_value.update538        with self.assertRaises(AirflowException) as cm:539            self.spanner_hook_no_default_project_id.update_database(540                instance_id=SPANNER_INSTANCE,541                database_id=SPANNER_DATABASE,542                ddl_statements=[])543        get_client.assert_not_called()544        instance_method.assert_not_called()545        database_method.assert_not_called()546        database_update_method.assert_not_called()547        err = cm.exception548        self.assertIn("The project id must be passed", str(err))549    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')550    def test_update_database_overridden_project_id(self, get_client):551        instance_method = get_client.return_value.instance552        instance_exists_method = instance_method.return_value.exists553        instance_exists_method.return_value = True554        database_method = instance_method.return_value.database555        database_update_ddl_method = database_method.return_value.update_ddl556        res = self.spanner_hook_no_default_project_id.update_database(557            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,558            instance_id=SPANNER_INSTANCE,559            database_id=SPANNER_DATABASE,560            ddl_statements=[])561        get_client.assert_called_once_with(project_id='example-project')562        instance_method.assert_called_once_with(instance_id='instance')563        database_method.assert_called_with(database_id='database-name')564        database_update_ddl_method.assert_called_with(ddl_statements=[], operation_id=None)565        self.assertIsNone(res)566    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')567    def test_update_database_overridden_project_id_and_operation(self, get_client):568        instance_method = get_client.return_value.instance569        instance_exists_method = instance_method.return_value.exists570        instance_exists_method.return_value = True571        database_method = instance_method.return_value.database572        database_update_ddl_method = database_method.return_value.update_ddl573        res = self.spanner_hook_no_default_project_id.update_database(574            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,575            instance_id=SPANNER_INSTANCE,576            database_id=SPANNER_DATABASE,577            operation_id="operation",578            ddl_statements=[])579        get_client.assert_called_once_with(project_id='example-project')580        instance_method.assert_called_once_with(instance_id='instance')581        database_method.assert_called_with(database_id='database-name')582        database_update_ddl_method.assert_called_with(ddl_statements=[], operation_id="operation")583        self.assertIsNone(res)584    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')585    def test_delete_database_missing_project_id(self, get_client):586        instance_method = get_client.return_value.instance587        instance_exists_method = instance_method.return_value.exists588        instance_exists_method.return_value = True589        database_method = instance_method.return_value.database590        database_drop_method = database_method.return_value.drop591        database_exists_method = database_method.return_value.exists592        database_exists_method.return_value = True593        with self.assertRaises(AirflowException) as cm:594            self.spanner_hook_no_default_project_id.delete_database(595                instance_id=SPANNER_INSTANCE,596                database_id=SPANNER_DATABASE)597        get_client.assert_not_called()598        instance_method.assert_not_called()599        database_method.assert_not_called()600        database_exists_method.assert_not_called()601        database_drop_method.assert_not_called()602        err = cm.exception603        self.assertIn("The project id must be passed", str(err))604    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')605    def test_delete_database_overridden_project_id(self, get_client):606        instance_method = get_client.return_value.instance607        instance_exists_method = instance_method.return_value.exists608        instance_exists_method.return_value = True609        database_method = instance_method.return_value.database610        database_drop_method = database_method.return_value.drop611        database_exists_method = database_method.return_value.exists612        database_exists_method.return_value = True613        res = self.spanner_hook_no_default_project_id.delete_database(614            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,615            instance_id=SPANNER_INSTANCE,616            database_id=SPANNER_DATABASE)617        get_client.assert_called_once_with(project_id='example-project')618        instance_method.assert_called_once_with(instance_id='instance')619        database_method.assert_called_with(database_id='database-name')620        database_exists_method.assert_called_with()621        database_drop_method.assert_called_with()622        self.assertIsNone(res)623    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')624    def test_delete_database_missing_database(self, get_client):625        instance_method = get_client.return_value.instance626        instance_exists_method = instance_method.return_value.exists627        instance_exists_method.return_value = True628        database_method = instance_method.return_value.database629        database_drop_method = database_method.return_value.drop630        database_exists_method = database_method.return_value.exists631        database_exists_method.return_value = False632        self.spanner_hook_no_default_project_id.delete_database(633            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,634            instance_id=SPANNER_INSTANCE,635            database_id=SPANNER_DATABASE)636        get_client.assert_called_once_with(project_id='example-project')637        instance_method.assert_called_once_with(instance_id='instance')638        database_method.assert_called_with(database_id='database-name')639        database_exists_method.assert_called_with()640        database_drop_method.assert_not_called()641    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')642    def test_execute_dml_missing_project_id(self, get_client):643        instance_method = get_client.return_value.instance644        instance_exists_method = instance_method.return_value.exists645        instance_exists_method.return_value = True646        database_method = instance_method.return_value.database647        run_in_transaction_method = database_method.return_value.run_in_transaction648        with self.assertRaises(AirflowException) as cm:649            self.spanner_hook_no_default_project_id.execute_dml(650                instance_id=SPANNER_INSTANCE,651                database_id=SPANNER_DATABASE,652                queries='')653        get_client.assert_not_called()654        instance_method.assert_not_called()655        database_method.assert_not_called()656        run_in_transaction_method.assert_not_called()657        err = cm.exception658        self.assertIn("The project id must be passed", str(err))659    @mock.patch('airflow.contrib.hooks.gcp_spanner_hook.CloudSpannerHook._get_client')660    def test_execute_dml_overridden_project_id(self, get_client):661        instance_method = get_client.return_value.instance662        instance_exists_method = instance_method.return_value.exists663        instance_exists_method.return_value = True664        database_method = instance_method.return_value.database665        run_in_transaction_method = database_method.return_value.run_in_transaction666        res = self.spanner_hook_no_default_project_id.execute_dml(667            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,668            instance_id=SPANNER_INSTANCE,669            database_id=SPANNER_DATABASE,670            queries='')671        get_client.assert_called_once_with(project_id='example-project')672        instance_method.assert_called_once_with(instance_id='instance')673        database_method.assert_called_with(database_id='database-name')674        run_in_transaction_method.assert_called_with(mock.ANY)...

test_gcp_bigtable_hook.py

Source:test_gcp_bigtable_hook.py

1# -*- coding: utf-8 -*-2#3# Licensed to the Apache Software Foundation (ASF) under one4# or more contributor license agreements.  See the NOTICE file5# distributed with this work for additional information6# regarding copyright ownership.  The ASF licenses this file7# to you under the Apache License, Version 2.0 (the8# "License"); you may not use this file except in compliance9# with the License.  You may obtain a copy of the License at10#11#   http://www.apache.org/licenses/LICENSE-2.012#13# Unless required by applicable law or agreed to in writing,14# software distributed under the License is distributed on an15# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY16# KIND, either express or implied.  See the License for the17# specific language governing permissions and limitations18# under the License.19import google20import unittest21from google.cloud.bigtable import Client22from google.cloud.bigtable.instance import Instance23from tests.contrib.utils.base_gcp_mock import mock_base_gcp_hook_no_default_project_id, \24    mock_base_gcp_hook_default_project_id, GCP_PROJECT_ID_HOOK_UNIT_TEST25from tests.compat import mock26from airflow import AirflowException27from airflow.contrib.hooks.gcp_bigtable_hook import BigtableHook28CBT_INSTANCE = 'instance'29CBT_CLUSTER = 'cluster'30CBT_ZONE = 'zone'31CBT_TABLE = 'table'32class TestBigtableHookNoDefaultProjectId(unittest.TestCase):33    def setUp(self):34        with mock.patch('airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__',35                        new=mock_base_gcp_hook_no_default_project_id):36            self.bigtable_hook_no_default_project_id = BigtableHook(gcp_conn_id='test')37    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')38    def test_get_instance_missing_project_id(self, get_client):39        instance_method = get_client.return_value.instance40        instance_exists_method = instance_method.return_value.exists41        instance_exists_method.return_value = True42        with self.assertRaises(AirflowException) as cm:43            self.bigtable_hook_no_default_project_id.get_instance(instance_id=CBT_INSTANCE)44        instance_exists_method.assert_not_called()45        instance_method.assert_not_called()46        err = cm.exception47        self.assertIn("The project id must be passed", str(err))48    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')49    def test_get_instance_overridden_project_id(self, get_client):50        instance_method = get_client.return_value.instance51        instance_exists_method = instance_method.return_value.exists52        instance_exists_method.return_value = True53        res = self.bigtable_hook_no_default_project_id.get_instance(54            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,55            instance_id=CBT_INSTANCE)56        instance_method.assert_called_once_with('instance')57        instance_exists_method.assert_called_once_with()58        get_client.assert_called_once_with(project_id='example-project')59        self.assertIsNotNone(res)60    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')61    def test_delete_instance_missing_project_id(self, get_client):62        instance_method = get_client.return_value.instance63        instance_exists_method = instance_method.return_value.exists64        delete_method = instance_method.return_value.delete65        instance_exists_method.return_value = True66        with self.assertRaises(AirflowException) as cm:67            self.bigtable_hook_no_default_project_id.delete_instance(instance_id=CBT_INSTANCE)68        instance_exists_method.assert_not_called()69        instance_method.assert_not_called()70        delete_method.assert_not_called()71        err = cm.exception72        self.assertIn("The project id must be passed", str(err))73    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')74    def test_delete_instance_overridden_project_id(self, get_client):75        instance_method = get_client.return_value.instance76        instance_exists_method = instance_method.return_value.exists77        instance_exists_method.return_value = True78        delete_method = instance_method.return_value.delete79        res = self.bigtable_hook_no_default_project_id.delete_instance(80            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST, instance_id=CBT_INSTANCE)81        instance_method.assert_called_once_with('instance')82        instance_exists_method.assert_called_once_with()83        delete_method.assert_called_once_with()84        get_client.assert_called_once_with(project_id='example-project')85        self.assertIsNone(res)86    @mock.patch('google.cloud.bigtable.instance.Instance.create')87    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')88    def test_create_instance_missing_project_id(self, get_client, instance_create):89        operation = mock.Mock()90        operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client)91        instance_create.return_value = operation92        with self.assertRaises(AirflowException) as cm:93            self.bigtable_hook_no_default_project_id.create_instance(94                instance_id=CBT_INSTANCE,95                main_cluster_id=CBT_CLUSTER,96                main_cluster_zone=CBT_ZONE)97        get_client.assert_not_called()98        instance_create.assert_not_called()99        err = cm.exception100        self.assertIn("The project id must be passed", str(err))101    @mock.patch('google.cloud.bigtable.instance.Instance.create')102    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')103    def test_create_instance_overridden_project_id(self, get_client, instance_create):104        operation = mock.Mock()105        operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client)106        instance_create.return_value = operation107        res = self.bigtable_hook_no_default_project_id.create_instance(108            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,109            instance_id=CBT_INSTANCE,110            main_cluster_id=CBT_CLUSTER,111            main_cluster_zone=CBT_ZONE)112        get_client.assert_called_once_with(project_id='example-project')113        instance_create.assert_called_once_with(clusters=mock.ANY)114        self.assertEqual(res.instance_id, 'instance')115    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')116    def test_delete_table_missing_project_id(self, get_client):117        instance_method = get_client.return_value.instance118        instance_exists_method = instance_method.return_value.exists119        table_delete_method = instance_method.return_value.table.return_value.delete120        instance_exists_method.return_value = True121        with self.assertRaises(AirflowException) as cm:122            self.bigtable_hook_no_default_project_id.delete_table(123                instance_id=CBT_INSTANCE,124                table_id=CBT_TABLE)125        get_client.assert_not_called()126        instance_exists_method.assert_not_called()127        table_delete_method.assert_not_called()128        err = cm.exception129        self.assertIn("The project id must be passed", str(err))130    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')131    def test_delete_table_overridden_project_id(self, get_client):132        instance_method = get_client.return_value.instance133        instance_exists_method = instance_method.return_value.exists134        table_delete_method = instance_method.return_value.table.return_value.delete135        instance_exists_method.return_value = True136        self.bigtable_hook_no_default_project_id.delete_table(137            project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,138            instance_id=CBT_INSTANCE,139            table_id=CBT_TABLE)140        get_client.assert_called_once_with(project_id='example-project')141        instance_exists_method.assert_called_once_with()142        table_delete_method.assert_called_once_with()143class TestBigtableHookDefaultProjectId(unittest.TestCase):144    def setUp(self):145        with mock.patch('airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__',146                        new=mock_base_gcp_hook_default_project_id):147            self.bigtable_hook_default_project_id = BigtableHook(gcp_conn_id='test')148    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')149    def test_get_instance(self, get_client):150        instance_method = get_client.return_value.instance151        instance_exists_method = instance_method.return_value.exists152        instance_exists_method.return_value = True153        res = self.bigtable_hook_default_project_id.get_instance(154            instance_id=CBT_INSTANCE)155        instance_method.assert_called_once_with('instance')156        instance_exists_method.assert_called_once_with()157        get_client.assert_called_once_with(project_id='example-project')158        self.assertIsNotNone(res)159    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')160    def test_get_instance_overridden_project_id(self, get_client):161        instance_method = get_client.return_value.instance162        instance_exists_method = instance_method.return_value.exists163        instance_exists_method.return_value = True164        res = self.bigtable_hook_default_project_id.get_instance(165            project_id='new-project',166            instance_id=CBT_INSTANCE)167        instance_method.assert_called_once_with('instance')168        instance_exists_method.assert_called_once_with()169        get_client.assert_called_once_with(project_id='new-project')170        self.assertIsNotNone(res)171    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')172    def test_get_instance_no_instance(self, get_client):173        instance_method = get_client.return_value.instance174        instance_exists_method = instance_method.return_value.exists175        instance_exists_method.return_value = False176        res = self.bigtable_hook_default_project_id.get_instance(177            instance_id=CBT_INSTANCE)178        instance_method.assert_called_once_with('instance')179        instance_exists_method.assert_called_once_with()180        get_client.assert_called_once_with(project_id='example-project')181        self.assertIsNone(res)182    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')183    def test_delete_instance(self, get_client):184        instance_method = get_client.return_value.instance185        instance_exists_method = instance_method.return_value.exists186        instance_exists_method.return_value = True187        delete_method = instance_method.return_value.delete188        res = self.bigtable_hook_default_project_id.delete_instance(189            instance_id=CBT_INSTANCE)190        instance_method.assert_called_once_with('instance')191        instance_exists_method.assert_called_once_with()192        delete_method.assert_called_once_with()193        get_client.assert_called_once_with(project_id='example-project')194        self.assertIsNone(res)195    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')196    def test_delete_instance_overridden_project_id(self, get_client):197        instance_method = get_client.return_value.instance198        instance_exists_method = instance_method.return_value.exists199        instance_exists_method.return_value = True200        delete_method = instance_method.return_value.delete201        res = self.bigtable_hook_default_project_id.delete_instance(202            project_id='new-project', instance_id=CBT_INSTANCE)203        instance_method.assert_called_once_with('instance')204        instance_exists_method.assert_called_once_with()205        delete_method.assert_called_once_with()206        get_client.assert_called_once_with(project_id='new-project')207        self.assertIsNone(res)208    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')209    def test_delete_instance_no_instance(self, get_client):210        instance_method = get_client.return_value.instance211        instance_exists_method = instance_method.return_value.exists212        instance_exists_method.return_value = False213        delete_method = instance_method.return_value.delete214        self.bigtable_hook_default_project_id.delete_instance(215            instance_id=CBT_INSTANCE)216        instance_method.assert_called_once_with('instance')217        instance_exists_method.assert_called_once_with()218        delete_method.assert_not_called()219        get_client.assert_called_once_with(project_id='example-project')220    @mock.patch('google.cloud.bigtable.instance.Instance.create')221    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')222    def test_create_instance(self, get_client, instance_create):223        operation = mock.Mock()224        operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client)225        instance_create.return_value = operation226        res = self.bigtable_hook_default_project_id.create_instance(227            instance_id=CBT_INSTANCE,228            main_cluster_id=CBT_CLUSTER,229            main_cluster_zone=CBT_ZONE)230        get_client.assert_called_once_with(project_id='example-project')231        instance_create.assert_called_once_with(clusters=mock.ANY)232        self.assertEqual(res.instance_id, 'instance')233    @mock.patch('google.cloud.bigtable.instance.Instance.create')234    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')235    def test_create_instance_overridden_project_id(self, get_client, instance_create):236        operation = mock.Mock()237        operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client)238        instance_create.return_value = operation239        res = self.bigtable_hook_default_project_id.create_instance(240            project_id='new-project',241            instance_id=CBT_INSTANCE,242            main_cluster_id=CBT_CLUSTER,243            main_cluster_zone=CBT_ZONE)244        get_client.assert_called_once_with(project_id='new-project')245        instance_create.assert_called_once_with(clusters=mock.ANY)246        self.assertEqual(res.instance_id, 'instance')247    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')248    def test_delete_table(self, get_client):249        instance_method = get_client.return_value.instance250        instance_exists_method = instance_method.return_value.exists251        table_delete_method = instance_method.return_value.table.return_value.delete252        instance_exists_method.return_value = True253        self.bigtable_hook_default_project_id.delete_table(254            instance_id=CBT_INSTANCE,255            table_id=CBT_TABLE)256        get_client.assert_called_once_with(project_id='example-project')257        instance_exists_method.assert_called_once_with()258        table_delete_method.assert_called_once_with()259    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')260    def test_delete_table_overridden_project_id(self, get_client):261        instance_method = get_client.return_value.instance262        instance_exists_method = instance_method.return_value.exists263        table_delete_method = instance_method.return_value.table.return_value.delete264        instance_exists_method.return_value = True265        self.bigtable_hook_default_project_id.delete_table(266            project_id='new-project',267            instance_id=CBT_INSTANCE,268            table_id=CBT_TABLE)269        get_client.assert_called_once_with(project_id='new-project')270        instance_exists_method.assert_called_once_with()271        table_delete_method.assert_called_once_with()272    @mock.patch('google.cloud.bigtable.table.Table.create')273    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')274    def test_create_table(self, get_client, create):275        instance_method = get_client.return_value.instance276        instance_exists_method = instance_method.return_value.exists277        instance_exists_method.return_value = True278        client = mock.Mock(Client)279        instance = google.cloud.bigtable.instance.Instance(280            instance_id=CBT_INSTANCE,281            client=client)282        self.bigtable_hook_default_project_id.create_table(283            instance=instance,284            table_id=CBT_TABLE)285        get_client.assert_not_called()286        create.assert_called_once_with([], {})287    @mock.patch('google.cloud.bigtable.cluster.Cluster.update')288    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')289    def test_update_cluster(self, get_client, update):290        instance_method = get_client.return_value.instance291        instance_exists_method = instance_method.return_value.exists292        instance_exists_method.return_value = True293        client = mock.Mock(Client)294        instance = google.cloud.bigtable.instance.Instance(295            instance_id=CBT_INSTANCE,296            client=client)297        self.bigtable_hook_default_project_id.update_cluster(298            instance=instance,299            cluster_id=CBT_CLUSTER,300            nodes=4)301        get_client.assert_not_called()302        update.assert_called_once_with()303    @mock.patch('google.cloud.bigtable.table.Table.list_column_families')304    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')305    def test_list_column_families(self, get_client, list_column_families):306        instance_method = get_client.return_value.instance307        instance_exists_method = instance_method.return_value.exists308        instance_exists_method.return_value = True309        client = mock.Mock(Client)310        get_client.return_value = client311        instance = google.cloud.bigtable.instance.Instance(312            instance_id=CBT_INSTANCE,313            client=client)314        self.bigtable_hook_default_project_id.get_column_families_for_table(315            instance=instance, table_id=CBT_TABLE)316        get_client.assert_not_called()317        list_column_families.assert_called_once_with()318    @mock.patch('google.cloud.bigtable.table.Table.get_cluster_states')319    @mock.patch('airflow.contrib.hooks.gcp_bigtable_hook.BigtableHook._get_client')320    def test_get_cluster_states(self, get_client, get_cluster_states):321        instance_method = get_client.return_value.instance322        instance_exists_method = instance_method.return_value.exists323        instance_exists_method.return_value = True324        client = mock.Mock(Client)325        instance = google.cloud.bigtable.instance.Instance(326            instance_id=CBT_INSTANCE,327            client=client)328        self.bigtable_hook_default_project_id.get_cluster_states_for_table(329            instance=instance, table_id=CBT_TABLE)330        get_client.assert_not_called()...

models.py

Source:models.py

...31    # introspection support:32    __members__ = property(lambda x: x.__dir__())33    # Need to pretend to be the wrapped class, for the sake of objects that care34    # about this (especially in equality tests)35    __class__ = property(lambda x: get_client().__class__)36    __dict__ = property(lambda o: get_client().__dict__)37    __repr__ = lambda x: repr(get_client())38    __getattr__ = lambda x, o: getattr(get_client(), o)39    __setattr__ = lambda x, o, v: setattr(get_client(), o, v)40    __delattr__ = lambda x, o: delattr(get_client(), o)41    __lt__ = lambda x, o: get_client() < o42    __le__ = lambda x, o: get_client() <= o43    __eq__ = lambda x, o: get_client() == o44    __ne__ = lambda x, o: get_client() != o45    __gt__ = lambda x, o: get_client() > o46    __ge__ = lambda x, o: get_client() >= o47    if PY2:48        __cmp__ = lambda x, o: cmp(get_client(), o)  # NOQA49    __hash__ = lambda x: hash(get_client())50    # attributes are currently not callable51    # __call__ = lambda x, *a, **kw: get_client()(*a, **kw)52    __nonzero__ = lambda x: bool(get_client())53    __len__ = lambda x: len(get_client())54    __getitem__ = lambda x, i: get_client()[i]55    __iter__ = lambda x: iter(get_client())56    __contains__ = lambda x, i: i in get_client()57    __getslice__ = lambda x, i, j: get_client()[i:j]58    __add__ = lambda x, o: get_client() + o59    __sub__ = lambda x, o: get_client() - o60    __mul__ = lambda x, o: get_client() * o61    __floordiv__ = lambda x, o: get_client() // o62    __mod__ = lambda x, o: get_client() % o63    __divmod__ = lambda x, o: get_client().__divmod__(o)64    __pow__ = lambda x, o: get_client() ** o65    __lshift__ = lambda x, o: get_client() << o66    __rshift__ = lambda x, o: get_client() >> o67    __and__ = lambda x, o: get_client() & o68    __xor__ = lambda x, o: get_client() ^ o69    __or__ = lambda x, o: get_client() | o70    __div__ = lambda x, o: get_client().__div__(o)71    __truediv__ = lambda x, o: get_client().__truediv__(o)72    __neg__ = lambda x: -(get_client())73    __pos__ = lambda x: +(get_client())74    __abs__ = lambda x: abs(get_client())75    __invert__ = lambda x: ~(get_client())76    __complex__ = lambda x: complex(get_client())77    __int__ = lambda x: int(get_client())78    if PY2:79        __long__ = lambda x: long(get_client())  # NOQA80    __float__ = lambda x: float(get_client())81    __str__ = lambda x: binary_type(get_client())82    __unicode__ = lambda x: text_type(get_client())83    __oct__ = lambda x: oct(get_client())84    __hex__ = lambda x: hex(get_client())85    __index__ = lambda x: get_client().__index__()86    __coerce__ = lambda x, o: x.__coerce__(x, o)87    __enter__ = lambda x: x.__enter__()88    __exit__ = lambda x, *a, **kw: x.__exit__(*a, **kw)89client = ProxyClient()90def get_client(client=None, reset=False):91    global _client92    tmp_client = client is not None93    if not tmp_client:94        client = getattr(settings, 'SENTRY_CLIENT', 'raven.contrib.django.DjangoClient')95    if _client[0] != client or reset:96        options = convert_options(97            settings,98            defaults={99                'include_paths': get_installed_apps(),100            },101        )102        try:103            Client = import_string(client)104        except ImportError:105            logger.exception('Failed to import client: %s', client)106            if not _client[1]:107                # If there is no previous client, set the default one.108                client = 'raven.contrib.django.DjangoClient'109                _client = (client, get_client(client))110        else:111            instance = Client(**options)112            if not tmp_client:113                _client = (client, instance)114            return instance115    return _client[1]116def sentry_exception_handler(request=None, **kwargs):117    try:118        client.captureException(exc_info=sys.exc_info(), request=request)119    except Exception as exc:120        try:121            logger.exception('Unable to process log entry: %s' % (exc,))122        except Exception as exc:123            warnings.warn('Unable to process log entry: %s' % (exc,))124class SentryDjangoHandler(object):125    def __init__(self, client=client):126        self.client = client127        try:128            import celery129        except ImportError:130            self.has_celery = False131        else:132            self.has_celery = celery.VERSION >= (2, 5)133        self.celery_handler = None134    def install_celery(self):135        from raven.contrib.celery import (136            SentryCeleryHandler, register_logger_signal137        )138        ignore_expected = getattr(settings,139                                  'SENTRY_CELERY_IGNORE_EXPECTED',140                                  False)141        self.celery_handler = SentryCeleryHandler(client,142                                                  ignore_expected=ignore_expected)\143                                                  .install()144        # try:145        #     ga = lambda x, d=None: getattr(settings, 'SENTRY_%s' % x, d)146        #     options = getattr(settings, 'RAVEN_CONFIG', {})147        #     loglevel = options.get('celery_loglevel',148        #                            ga('CELERY_LOGLEVEL', logging.ERROR))149        #     register_logger_signal(client, loglevel=loglevel)150        # except Exception:151        #     logger.exception('Failed to install Celery error handler')152    def install(self):153        request_started.connect(self.before_request, weak=False)154        got_request_exception.connect(self.exception_handler, weak=False)155        if self.has_celery:156            try:157                self.install_celery()158            except Exception:159                logger.exception('Failed to install Celery error handler')160    def uninstall(self):161        request_started.disconnect(self.before_request)162        got_request_exception.disconnect(self.exception_handler)163        if self.celery_handler:164            self.celery_handler.uninstall()165    def exception_handler(self, request=None, **kwargs):166        try:167            self.client.captureException(exc_info=sys.exc_info(), request=request)168        except Exception as exc:169            try:170                logger.exception('Unable to process log entry: %s' % (exc,))171            except Exception as exc:172                warnings.warn('Unable to process log entry: %s' % (exc,))173    def before_request(self, *args, **kwargs):174        self.client.context.activate()175def register_serializers():176    # force import so serializers can call register177    import raven.contrib.django.serializers  # NOQA178def install_middleware(middleware_name, lookup_names=None):179    """180    Install specified middleware181    """182    if lookup_names is None:183        lookup_names = (middleware_name,)184    # default settings.MIDDLEWARE is None185    middleware_attr = 'MIDDLEWARE' if getattr(settings,186                                              'MIDDLEWARE',187                                              None) is not None \188        else 'MIDDLEWARE_CLASSES'189    # make sure to get an empty tuple when attr is None190    middleware = getattr(settings, middleware_attr, ()) or ()191    if set(lookup_names).isdisjoint(set(middleware)):192        setattr(settings,193                middleware_attr,194                type(middleware)((middleware_name,)) + middleware)195_setup_lock = Lock()196_initialized = False197def initialize():198    global _initialized199    with _setup_lock:200        if _initialized:201            return202        # mark this as initialized immediatley to avoid recursive import issues203        _initialized = True204        try:205            register_serializers()206            install_middleware(207                'raven.contrib.django.middleware.SentryMiddleware',208                (209                    'raven.contrib.django.middleware.SentryMiddleware',210                    'raven.contrib.django.middleware.SentryLogMiddleware'))211            install_middleware(212                'raven.contrib.django.middleware.DjangoRestFrameworkCompatMiddleware')213            # XXX(dcramer): maybe this setting should disable ALL of this?214            if not getattr(settings, 'DISABLE_SENTRY_INSTRUMENTATION', False):215                handler = SentryDjangoHandler()216                handler.install()217            # instantiate client so hooks get registered218            get_client()  # NOQA219        except Exception:220            _initialized = False221# Django 1.7 uses ``raven.contrib.django.apps.RavenConfig``222if django.VERSION < (1, 7, 0):...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.