Skip to content

Commit 810c15e

Browse files
author
Tobiasz Kędzierski
authored
Fix and improve GCP BigTable hook and system test (#13896)
Improve environment variables in GCP BigTable system test. It will help to parametrize system tests.
1 parent 6616617 commit 810c15e

4 files changed

Lines changed: 81 additions & 25 deletions

File tree

‎airflow/providers/google/cloud/example_dags/example_bigtable.py‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -60,22 +60,22 @@
6060
from airflow.utils.dates import days_ago
6161

6262
GCP_PROJECT_ID = getenv('GCP_PROJECT_ID', 'example-project')
63-
CBT_INSTANCE_ID = getenv('CBT_INSTANCE_ID', 'some-instance-id')
64-
CBT_INSTANCE_DISPLAY_NAME = getenv('CBT_INSTANCE_DISPLAY_NAME', 'Human-readable name')
63+
CBT_INSTANCE_ID = getenv('GCP_BIG_TABLE_INSTANCE_ID', 'some-instance-id')
64+
CBT_INSTANCE_DISPLAY_NAME = getenv('GCP_BIG_TABLE_INSTANCE_DISPLAY_NAME', 'Human-readable name')
6565
CBT_INSTANCE_DISPLAY_NAME_UPDATED = getenv(
66-
"CBT_INSTANCE_DISPLAY_NAME_UPDATED", "Human-readable name - updated"
66+
"GCP_BIG_TABLE_INSTANCE_DISPLAY_NAME_UPDATED", f"{CBT_INSTANCE_DISPLAY_NAME} - updated"
6767
)
68-
CBT_INSTANCE_TYPE = getenv('CBT_INSTANCE_TYPE', '2')
69-
CBT_INSTANCE_TYPE_PROD = getenv('CBT_INSTANCE_TYPE_PROD', '1')
70-
CBT_INSTANCE_LABELS = getenv('CBT_INSTANCE_LABELS', '{}')
71-
CBT_INSTANCE_LABELS_UPDATED = getenv('CBT_INSTANCE_LABELS', '{"env": "prod"}')
72-
CBT_CLUSTER_ID = getenv('CBT_CLUSTER_ID', 'some-cluster-id')
73-
CBT_CLUSTER_ZONE = getenv('CBT_CLUSTER_ZONE', 'europe-west1-b')
74-
CBT_CLUSTER_NODES = getenv('CBT_CLUSTER_NODES', '3')
75-
CBT_CLUSTER_NODES_UPDATED = getenv('CBT_CLUSTER_NODES_UPDATED', '5')
76-
CBT_CLUSTER_STORAGE_TYPE = getenv('CBT_CLUSTER_STORAGE_TYPE', '2')
77-
CBT_TABLE_ID = getenv('CBT_TABLE_ID', 'some-table-id')
78-
CBT_POKE_INTERVAL = getenv('CBT_POKE_INTERVAL', '60')
68+
CBT_INSTANCE_TYPE = getenv('GCP_BIG_TABLE_INSTANCE_TYPE', '2')
69+
CBT_INSTANCE_TYPE_PROD = getenv('GCP_BIG_TABLE_INSTANCE_TYPE_PROD', '1')
70+
CBT_INSTANCE_LABELS = getenv('GCP_BIG_TABLE_INSTANCE_LABELS', '{}')
71+
CBT_INSTANCE_LABELS_UPDATED = getenv('GCP_BIG_TABLE_INSTANCE_LABELS_UPDATED', '{"env": "prod"}')
72+
CBT_CLUSTER_ID = getenv('GCP_BIG_TABLE_CLUSTER_ID', 'some-cluster-id')
73+
CBT_CLUSTER_ZONE = getenv('GCP_BIG_TABLE_CLUSTER_ZONE', 'europe-west1-b')
74+
CBT_CLUSTER_NODES = getenv('GCP_BIG_TABLE_CLUSTER_NODES', '3')
75+
CBT_CLUSTER_NODES_UPDATED = getenv('GCP_BIG_TABLE_CLUSTER_NODES_UPDATED', '5')
76+
CBT_CLUSTER_STORAGE_TYPE = getenv('GCP_BIG_TABLE_CLUSTER_STORAGE_TYPE', '2')
77+
CBT_TABLE_ID = getenv('GCP_BIG_TABLE_TABLE_ID', 'some-table-id')
78+
CBT_POKE_INTERVAL = getenv('GCP_BIG_TABLE_POKE_INTERVAL', '60')
7979

8080

8181
with models.DAG(
@@ -93,8 +93,8 @@
9393
instance_display_name=CBT_INSTANCE_DISPLAY_NAME,
9494
instance_type=int(CBT_INSTANCE_TYPE),
9595
instance_labels=json.loads(CBT_INSTANCE_LABELS),
96-
cluster_nodes=int(CBT_CLUSTER_NODES),
97-
cluster_storage_type=CBT_CLUSTER_STORAGE_TYPE,
96+
cluster_nodes=None,
97+
cluster_storage_type=int(CBT_CLUSTER_STORAGE_TYPE),
9898
task_id='create_instance_task',
9999
)
100100
create_instance_task2 = BigtableCreateInstanceOperator(

‎airflow/providers/google/cloud/hooks/bigtable.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -169,7 +169,14 @@ def create_instance(
169169
instance_labels,
170170
)
171171

172-
clusters = [instance.cluster(main_cluster_id, main_cluster_zone, cluster_nodes, cluster_storage_type)]
172+
cluster_kwargs = dict(
173+
cluster_id=main_cluster_id,
174+
location_id=main_cluster_zone,
175+
default_storage_type=cluster_storage_type,
176+
)
177+
if instance_type != enums.Instance.Type.DEVELOPMENT and cluster_nodes:
178+
cluster_kwargs["serve_nodes"] = cluster_nodes
179+
clusters = [instance.cluster(**cluster_kwargs)]
173180
if replica_cluster_id and replica_cluster_zone:
174181
warnings.warn(
175182
"The replica_cluster_id and replica_cluster_zone parameter have been deprecated."

‎tests/providers/google/cloud/hooks/test_bigtable.py‎

Lines changed: 55 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -309,7 +309,7 @@ def test_create_instance(self, get_client, instance_create, mock_project_id):
309309
@mock.patch('google.cloud.bigtable.instance.Instance.cluster')
310310
@mock.patch('google.cloud.bigtable.instance.Instance.create')
311311
@mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client')
312-
def test_create_instance_with_one_replica_cluster(
312+
def test_create_instance_with_one_replica_cluster_production(
313313
self, get_client, instance_create, cluster, mock_project_id
314314
):
315315
operation = mock.Mock()
@@ -325,10 +325,57 @@ def test_create_instance_with_one_replica_cluster(
325325
cluster_nodes=1,
326326
cluster_storage_type=enums.StorageType.SSD,
327327
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,
328+
instance_type=enums.Instance.Type.PRODUCTION,
328329
)
329330
cluster.assert_has_calls(
330331
[
331-
unittest.mock.call(CBT_CLUSTER, CBT_ZONE, 1, enums.StorageType.SSD),
332+
unittest.mock.call(
333+
cluster_id=CBT_CLUSTER,
334+
location_id=CBT_ZONE,
335+
serve_nodes=1,
336+
default_storage_type=enums.StorageType.SSD,
337+
),
338+
unittest.mock.call(
339+
CBT_REPLICA_CLUSTER_ID, CBT_REPLICA_CLUSTER_ZONE, 1, enums.StorageType.SSD
340+
),
341+
],
342+
any_order=True,
343+
)
344+
get_client.assert_called_once_with(project_id='example-project')
345+
instance_create.assert_called_once_with(clusters=mock.ANY)
346+
assert res.instance_id == 'instance'
347+
348+
@mock.patch(
349+
'airflow.providers.google.common.hooks.base_google.GoogleBaseHook.project_id',
350+
new_callable=PropertyMock,
351+
return_value=GCP_PROJECT_ID_HOOK_UNIT_TEST,
352+
)
353+
@mock.patch('google.cloud.bigtable.instance.Instance.cluster')
354+
@mock.patch('google.cloud.bigtable.instance.Instance.create')
355+
@mock.patch('airflow.providers.google.cloud.hooks.bigtable.BigtableHook._get_client')
356+
def test_create_instance_with_one_replica_cluster_development(
357+
self, get_client, instance_create, cluster, mock_project_id
358+
):
359+
operation = mock.Mock()
360+
operation.result_return_value = Instance(instance_id=CBT_INSTANCE, client=get_client)
361+
instance_create.return_value = operation
362+
363+
res = self.bigtable_hook_default_project_id.create_instance(
364+
instance_id=CBT_INSTANCE,
365+
main_cluster_id=CBT_CLUSTER,
366+
main_cluster_zone=CBT_ZONE,
367+
replica_cluster_id=CBT_REPLICA_CLUSTER_ID,
368+
replica_cluster_zone=CBT_REPLICA_CLUSTER_ZONE,
369+
cluster_nodes=1,
370+
cluster_storage_type=enums.StorageType.SSD,
371+
project_id=GCP_PROJECT_ID_HOOK_UNIT_TEST,
372+
instance_type=enums.Instance.Type.DEVELOPMENT,
373+
)
374+
cluster.assert_has_calls(
375+
[
376+
unittest.mock.call(
377+
cluster_id=CBT_CLUSTER, location_id=CBT_ZONE, default_storage_type=enums.StorageType.SSD
378+
),
332379
unittest.mock.call(
333380
CBT_REPLICA_CLUSTER_ID, CBT_REPLICA_CLUSTER_ZONE, 1, enums.StorageType.SSD
334381
),
@@ -365,7 +412,12 @@ def test_create_instance_with_multiple_replica_clusters(
365412
)
366413
cluster.assert_has_calls(
367414
[
368-
unittest.mock.call(CBT_CLUSTER, CBT_ZONE, 1, enums.StorageType.SSD),
415+
unittest.mock.call(
416+
cluster_id=CBT_CLUSTER,
417+
location_id=CBT_ZONE,
418+
serve_nodes=1,
419+
default_storage_type=enums.StorageType.SSD,
420+
),
369421
unittest.mock.call('replica-1', 'us-west1-a', 1, enums.StorageType.SSD),
370422
unittest.mock.call('replica-2', 'us-central1-f', 1, enums.StorageType.SSD),
371423
unittest.mock.call('replica-3', 'us-east1-d', 1, enums.StorageType.SSD),

‎tests/providers/google/cloud/operators/test_bigtable_system.py‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,13 @@
1515
# KIND, either express or implied. See the License for the
1616
# specific language governing permissions and limitations
1717
# under the License.
18-
import os
1918

2019
import pytest
2120

21+
from airflow.providers.google.cloud.example_dags.example_bigtable import CBT_INSTANCE_ID, GCP_PROJECT_ID
2222
from tests.providers.google.cloud.utils.gcp_authenticator import GCP_BIGTABLE_KEY
2323
from tests.test_utils.gcp_system_helpers import CLOUD_DAG_FOLDER, GoogleSystemTest, provide_gcp_context
2424

25-
GCP_PROJECT_ID = os.environ.get('GCP_PROJECT_ID', 'example-project')
26-
CBT_INSTANCE = os.environ.get('CBT_INSTANCE_ID', 'testinstance')
27-
2825

2926
@pytest.mark.backend("mysql", "postgres")
3027
@pytest.mark.credential_file(GCP_BIGTABLE_KEY)
@@ -45,7 +42,7 @@ def tearDown(self):
4542
'--verbosity=none',
4643
'instances',
4744
'delete',
48-
CBT_INSTANCE,
45+
CBT_INSTANCE_ID,
4946
],
5047
key=GCP_BIGTABLE_KEY,
5148
)

0 commit comments

Comments
 (0)