Improve environment variables in GCP Memorystore system test (#13833)

apache · Jan 22, 2021 · 61c1d6e · 61c1d6e
1 parent 202f660
commit 61c1d6e
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 18 deletions.
diff --git a/airflow/providers/google/cloud/example_dags/example_cloud_memorystore.py b/airflow/providers/google/cloud/example_dags/example_cloud_memorystore.py
@@ -19,7 +19,6 @@
 Example Airflow DAG for Google Cloud Memorystore service.
 """
 import os
-from urllib.parse import urlparse
 
 from google.cloud.memcache_v1beta2.types import cloud_memcache
 from google.cloud.redis_v1 import FailoverInstanceRequest, Instance
@@ -51,20 +50,21 @@
 
 GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "example-project")
 
-MEMORYSTORE_REDIS_INSTANCE_NAME = os.environ.get("GCP_MEMORYSTORE_INSTANCE_NAME", "test-memorystoreredis-")
+MEMORYSTORE_REDIS_INSTANCE_NAME = os.environ.get(
+ "GCP_MEMORYSTORE_REDIS_INSTANCE_NAME", "test-memorystore-redis"
+)
 MEMORYSTORE_REDIS_INSTANCE_NAME_2 = os.environ.get(
- "GCP_MEMORYSTORE_INSTANCE_NAME2", "test-memorystore-redis-2"
+ "GCP_MEMORYSTORE_REDIS_INSTANCE_NAME_2", "test-memorystore-redis-2"
 )
 MEMORYSTORE_REDIS_INSTANCE_NAME_3 = os.environ.get(
- "GCP_MEMORYSTORE_INSTANCE_NAME3", "test-memorystore-redis-3"
+ "GCP_MEMORYSTORE_REDIS_INSTANCE_NAME_3", "test-memorystore-redis-3"
 )
 MEMORYSTORE_MEMCACHED_INSTANCE_NAME = os.environ.get(
- "GCP_MEMORYSTORE_INSTANCE_NAME4", "test-memorystore-memcached-1"
+ "GCP_MEMORYSTORE_MEMCACHED_INSTANCE_NAME", "test-memorystore-memcached-1"
 )
 
-EXPORT_GCS_URL = os.environ.get("GCP_MEMORYSTORE_EXPORT_GCS_URL", "gs://test-memorystore/my-export.rdb")
-EXPORT_GCS_URL_PARTS = urlparse(EXPORT_GCS_URL)
-BUCKET_NAME = EXPORT_GCS_URL_PARTS.netloc
+BUCKET_NAME = os.environ.get("GCP_MEMORYSTORE_BUCKET", "test-memorystore-bucket")
+EXPORT_GCS_URL = f"gs://{BUCKET_NAME}/my-export.rdb"
 
 # [START howto_operator_instance]
 FIRST_INSTANCE = {"tier": Instance.Tier.BASIC, "memory_size_gb": 1}

diff --git a/tests/providers/google/cloud/operators/test_cloud_memorystore_system.py b/tests/providers/google/cloud/operators/test_cloud_memorystore_system.py
@@ -16,19 +16,12 @@
 # specific language governing permissions and limitations
 # under the License.
 """System tests for Google Cloud Memorystore operators"""
-import os
-from urllib.parse import urlparse
-
 import pytest
 
+from airflow.providers.google.cloud.example_dags.example_cloud_memorystore import BUCKET_NAME
 from tests.providers.google.cloud.utils.gcp_authenticator import GCP_MEMORYSTORE
 from tests.test_utils.gcp_system_helpers import CLOUD_DAG_FOLDER, GoogleSystemTest, provide_gcp_context
 
-GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID", "example-project")
-GCP_ARCHIVE_URL = os.environ.get("GCP_MEMORYSTORE_EXPORT_GCS_URL", "gs://test-memorystore/my-export.rdb")
-GCP_ARCHIVE_URL_PARTS = urlparse(GCP_ARCHIVE_URL)
-GCP_BUCKET_NAME = GCP_ARCHIVE_URL_PARTS.netloc
-
 
 @pytest.mark.backend("mysql", "postgres")
 @pytest.mark.credential_file(GCP_MEMORYSTORE)
@@ -42,7 +35,7 @@ class CloudMemorystoreSystemTest(GoogleSystemTest):
  @provide_gcp_context(GCP_MEMORYSTORE)
  def setUp(self):
  super().setUp()
- self.create_gcs_bucket(GCP_BUCKET_NAME, location="europe-north1")
+ self.create_gcs_bucket(BUCKET_NAME, location="europe-north1")
 
  @provide_gcp_context(GCP_MEMORYSTORE)
  def test_run_example_dag_memorystore_redis(self):
@@ -54,5 +47,5 @@ def test_run_example_dag_memorystore_memcached(self):
 
  @provide_gcp_context(GCP_MEMORYSTORE)
  def tearDown(self):
- self.delete_gcs_bucket(GCP_BUCKET_NAME)
+ self.delete_gcs_bucket(BUCKET_NAME)
  super().tearDown()