From 25805db92786478d959ee3e2a1943e584e75c304 Mon Sep 17 00:00:00 2001
From: Ritesh Ghorse <riteshghorse@gmail.com>
Date: Thu, 28 Mar 2024 20:16:16 -0400
Subject: [PATCH] [Python] Add redis client to python dependencies (#30677)

* add redis client to python  dependencies

* fix lint

* update dependencies

* add upper bound

* remove redis requirements file and uses_redis tag

* new task for testcontainers

* fix expected value after syncing

* remove requirement file
---
 .../trigger_files/beam_PostCommit_Python.json  |  0
 .../transforms/elementwise/enrichment_test.py  |  6 +++---
 sdks/python/apache_beam/io/requestresponse.py  |  2 +-
 .../apache_beam/io/requestresponse_it_test.py  |  2 +-
 .../io/requestresponse_tests_requirements.txt  | 18 ------------------
 .../enrichment_handlers/bigtable_it_test.py    |  2 +-
 .../vertex_ai_feature_store_it_test.py         |  2 +-
 ...rtex_ai_feature_store_tests_requirement.txt | 18 ------------------
 .../py310/base_image_requirements.txt          |  4 +++-
 .../py311/base_image_requirements.txt          |  3 ++-
 .../container/py38/base_image_requirements.txt |  4 +++-
 .../container/py39/base_image_requirements.txt |  4 +++-
 sdks/python/pytest.ini                         |  2 +-
 sdks/python/setup.py                           |  1 +
 sdks/python/test-suites/dataflow/common.gradle |  9 +--------
 sdks/python/test-suites/direct/common.gradle   | 16 +++++-----------
 16 files changed, 26 insertions(+), 67 deletions(-)
 create mode 100644 .github/trigger_files/beam_PostCommit_Python.json
 delete mode 100644 sdks/python/apache_beam/io/requestresponse_tests_requirements.txt
 delete mode 100644 sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_tests_requirement.txt

diff --git a/.github/trigger_files/beam_PostCommit_Python.json b/.github/trigger_files/beam_PostCommit_Python.json
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py
index 3fd759bbc320..8a7cdfbe9263 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py
@@ -44,9 +44,9 @@ def validate_enrichment_with_bigtable():
 
 def validate_enrichment_with_vertex_ai():
   expected = '''[START enrichment_with_vertex_ai]
-Row(user_id='2963', product_id=14235, sale_price=15.0, age=29.0, gender='1', state='97', country='2')
-Row(user_id='21422', product_id=11203, sale_price=12.0, age=36.0, state='184', gender='1', country='5')
-Row(user_id='20592', product_id=8579, sale_price=9.0, age=30.0, state='86', gender='1', country='4')
+Row(user_id='2963', product_id=14235, sale_price=15.0, age=12.0, state='1', gender='1', country='1')
+Row(user_id='21422', product_id=11203, sale_price=12.0, age=12.0, state='0', gender='0', country='0')
+Row(user_id='20592', product_id=8579, sale_price=9.0, age=12.0, state='2', gender='1', country='2')
   [END enrichment_with_vertex_ai]'''.splitlines()[1:-1]
   return expected
 
diff --git a/sdks/python/apache_beam/io/requestresponse.py b/sdks/python/apache_beam/io/requestresponse.py
index 706bce95f5ee..4458aa59c18c 100644
--- a/sdks/python/apache_beam/io/requestresponse.py
+++ b/sdks/python/apache_beam/io/requestresponse.py
@@ -33,10 +33,10 @@
 from typing import TypeVar
 from typing import Union
 
+import redis
 from google.api_core.exceptions import TooManyRequests
 
 import apache_beam as beam
-import redis
 from apache_beam import pvalue
 from apache_beam.coders import coders
 from apache_beam.io.components.adaptive_throttler import AdaptiveThrottler
diff --git a/sdks/python/apache_beam/io/requestresponse_it_test.py b/sdks/python/apache_beam/io/requestresponse_it_test.py
index bd8c63dea587..712ccc7881d6 100644
--- a/sdks/python/apache_beam/io/requestresponse_it_test.py
+++ b/sdks/python/apache_beam/io/requestresponse_it_test.py
@@ -203,7 +203,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
     pass
 
 
-@pytest.mark.uses_redis
+@pytest.mark.uses_testcontainer
 class TestRedisCache(unittest.TestCase):
   def setUp(self) -> None:
     self.retries = 3
diff --git a/sdks/python/apache_beam/io/requestresponse_tests_requirements.txt b/sdks/python/apache_beam/io/requestresponse_tests_requirements.txt
deleted file mode 100644
index 1d8869705097..000000000000
--- a/sdks/python/apache_beam/io/requestresponse_tests_requirements.txt
+++ /dev/null
@@ -1,18 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-
-redis>=5.0.0
diff --git a/sdks/python/apache_beam/transforms/enrichment_handlers/bigtable_it_test.py b/sdks/python/apache_beam/transforms/enrichment_handlers/bigtable_it_test.py
index b792bc8ba946..bbabbc306e61 100644
--- a/sdks/python/apache_beam/transforms/enrichment_handlers/bigtable_it_test.py
+++ b/sdks/python/apache_beam/transforms/enrichment_handlers/bigtable_it_test.py
@@ -148,7 +148,7 @@ def create_rows(table):
     row.commit()
 
 
-@pytest.mark.uses_redis
+@pytest.mark.uses_testcontainer
 class TestBigTableEnrichment(unittest.TestCase):
   def setUp(self):
     self.project_id = 'apache-beam-testing'
diff --git a/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_it_test.py b/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_it_test.py
index 066553decf20..c5482309a251 100644
--- a/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_it_test.py
+++ b/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_it_test.py
@@ -62,7 +62,7 @@ def process(self, element: beam.Row, *args, **kwargs):
             f"from feature store")
 
 
-@pytest.mark.uses_redis
+@pytest.mark.uses_testcontainer
 class TestVertexAIFeatureStoreHandler(unittest.TestCase):
   def setUp(self) -> None:
     self.project = 'apache-beam-testing'
diff --git a/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_tests_requirement.txt b/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_tests_requirement.txt
deleted file mode 100644
index cd74683a51c1..000000000000
--- a/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_tests_requirement.txt
+++ /dev/null
@@ -1,18 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-redis
-google-cloud-aiplatform
diff --git a/sdks/python/container/py310/base_image_requirements.txt b/sdks/python/container/py310/base_image_requirements.txt
index 1258575696db..06fd375a8315 100644
--- a/sdks/python/container/py310/base_image_requirements.txt
+++ b/sdks/python/container/py310/base_image_requirements.txt
@@ -21,6 +21,7 @@
 # https://s.apache.org/beam-python-dev-wiki
 # Reach out to a committer if you need help.
 
+async-timeout==4.0.3
 attrs==23.2.0
 beautifulsoup4==4.12.3
 bs4==0.0.2
@@ -81,7 +82,7 @@ grpcio-status==1.62.1
 guppy3==3.1.4.post1
 hdfs==2.7.3
 httplib2==0.22.0
-hypothesis==6.99.12
+hypothesis==6.99.13
 idna==3.6
 iniconfig==2.0.0
 joblib==1.3.2
@@ -124,6 +125,7 @@ python-dateutil==2.9.0.post0
 python-snappy==0.7.1
 pytz==2024.1
 PyYAML==6.0.1
+redis==5.0.3
 referencing==0.34.0
 regex==2023.12.25
 requests==2.31.0
diff --git a/sdks/python/container/py311/base_image_requirements.txt b/sdks/python/container/py311/base_image_requirements.txt
index 5c703f727b93..71b3c48a35ba 100644
--- a/sdks/python/container/py311/base_image_requirements.txt
+++ b/sdks/python/container/py311/base_image_requirements.txt
@@ -79,7 +79,7 @@ grpcio-status==1.62.1
 guppy3==3.1.4.post1
 hdfs==2.7.3
 httplib2==0.22.0
-hypothesis==6.99.12
+hypothesis==6.99.13
 idna==3.6
 iniconfig==2.0.0
 joblib==1.3.2
@@ -121,6 +121,7 @@ pytest-xdist==3.5.0
 python-dateutil==2.9.0.post0
 pytz==2024.1
 PyYAML==6.0.1
+redis==5.0.3
 referencing==0.34.0
 regex==2023.12.25
 requests==2.31.0
diff --git a/sdks/python/container/py38/base_image_requirements.txt b/sdks/python/container/py38/base_image_requirements.txt
index 9a6e3af79536..887aa8a2531b 100644
--- a/sdks/python/container/py38/base_image_requirements.txt
+++ b/sdks/python/container/py38/base_image_requirements.txt
@@ -21,6 +21,7 @@
 # https://s.apache.org/beam-python-dev-wiki
 # Reach out to a committer if you need help.
 
+async-timeout==4.0.3
 attrs==23.2.0
 backports.zoneinfo==0.2.1
 beautifulsoup4==4.12.3
@@ -82,7 +83,7 @@ grpcio-status==1.62.1
 guppy3==3.1.4.post1
 hdfs==2.7.3
 httplib2==0.22.0
-hypothesis==6.99.12
+hypothesis==6.99.13
 idna==3.6
 importlib_metadata==7.1.0
 importlib_resources==6.4.0
@@ -128,6 +129,7 @@ python-dateutil==2.9.0.post0
 python-snappy==0.7.1
 pytz==2024.1
 PyYAML==6.0.1
+redis==5.0.3
 referencing==0.34.0
 regex==2023.12.25
 requests==2.31.0
diff --git a/sdks/python/container/py39/base_image_requirements.txt b/sdks/python/container/py39/base_image_requirements.txt
index cf116324eace..9d8009c0048f 100644
--- a/sdks/python/container/py39/base_image_requirements.txt
+++ b/sdks/python/container/py39/base_image_requirements.txt
@@ -21,6 +21,7 @@
 # https://s.apache.org/beam-python-dev-wiki
 # Reach out to a committer if you need help.
 
+async-timeout==4.0.3
 attrs==23.2.0
 beautifulsoup4==4.12.3
 bs4==0.0.2
@@ -81,7 +82,7 @@ grpcio-status==1.62.1
 guppy3==3.1.4.post1
 hdfs==2.7.3
 httplib2==0.22.0
-hypothesis==6.99.12
+hypothesis==6.99.13
 idna==3.6
 importlib_metadata==7.1.0
 iniconfig==2.0.0
@@ -125,6 +126,7 @@ python-dateutil==2.9.0.post0
 python-snappy==0.7.1
 pytz==2024.1
 PyYAML==6.0.1
+redis==5.0.3
 referencing==0.34.0
 regex==2023.12.25
 requests==2.31.0
diff --git a/sdks/python/pytest.ini b/sdks/python/pytest.ini
index e78697169bb0..8df74adbc62e 100644
--- a/sdks/python/pytest.ini
+++ b/sdks/python/pytest.ini
@@ -65,7 +65,7 @@ markers =
     uses_tf: tests that utilize tensorflow.
     uses_transformers: tests that utilize transformers in some way.
     vertex_ai_postcommit: vertex ai postcommits that need additional deps.
-    uses_redis: enrichment transform tests that need redis.
+    uses_testcontainer: tests that use testcontainers.
     uses_mock_api: tests that uses the mock API cluster.
 
 # Default timeout intended for unit tests.
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
index 409951cbc415..0e49f913fb44 100644
--- a/sdks/python/setup.py
+++ b/sdks/python/setup.py
@@ -392,6 +392,7 @@ def get_portability_package_data():
           'pydot>=1.2.0,<2',
           'python-dateutil>=2.8.0,<3',
           'pytz>=2018.3',
+          'redis>=5.0.0,<6',
           'regex>=2020.6.8',
           'requests>=2.24.0,<3.0.0',
           'typing-extensions>=3.7.0',
diff --git a/sdks/python/test-suites/dataflow/common.gradle b/sdks/python/test-suites/dataflow/common.gradle
index cadf3a6ae2c6..8f3402035eb2 100644
--- a/sdks/python/test-suites/dataflow/common.gradle
+++ b/sdks/python/test-suites/dataflow/common.gradle
@@ -491,13 +491,7 @@ task tftTests {
 task mockAPITests {
   dependsOn 'initializeForDataflowJob'
   dependsOn ':sdks:python:sdist'
-  def requirementsFile = "${rootDir}/sdks/python/apache_beam/io/requestresponse_tests_requirements.txt"
-  doFirst {
-      exec {
-        executable 'sh'
-        args '-c', ". ${envdir}/bin/activate && pip install -r $requirementsFile"
-      }
-    }
+
  doLast {
   def testOpts = basicTestOpts
   def argMap = [
@@ -506,7 +500,6 @@ task mockAPITests {
     "runner": "TestDataflowRunner",
     "project": "apache-beam-testing",
     "region": "us-west1",
-    "requirements_file": "$requirementsFile"
   ]
   def cmdArgs = mapToArgString(argMap)
   exec {
diff --git a/sdks/python/test-suites/direct/common.gradle b/sdks/python/test-suites/direct/common.gradle
index 084c51c9b4e4..1851a4d9cd0b 100644
--- a/sdks/python/test-suites/direct/common.gradle
+++ b/sdks/python/test-suites/direct/common.gradle
@@ -370,23 +370,17 @@ task transformersInferenceTest {
     }
 }
 
-// Enrichment transform tests that uses Redis
-task enrichmentRedisTest {
+// Integration tests that use testcontainers
+task testcontainersTest {
   dependsOn 'installGcpTest'
   dependsOn ':sdks:python:sdist'
-  def requirementsFile = "${rootDir}/sdks/python/apache_beam/transforms/enrichment_handlers/vertex_ai_feature_store_tests_requirement.txt"
-  doFirst {
-      exec {
-        executable 'sh'
-        args '-c', ". ${envdir}/bin/activate && pip install -r $requirementsFile"
-      }
-    }
+
   doLast {
       def testOpts = basicTestOpts
       def argMap = [
           "test_opts": testOpts,
           "suite": "postCommitIT-direct-py${pythonVersionSuffix}",
-          "collect": "uses_redis",
+          "collect": "uses_testcontainer",
           "runner": "TestDirectRunner",
           "region": "us-central1",
       ]
@@ -406,7 +400,7 @@ project.tasks.register("inferencePostCommitIT") {
   'tensorflowInferenceTest',
   'xgboostInferenceTest',
   'transformersInferenceTest',
-  'enrichmentRedisTest',
+  'testcontainersTest',
   // (TODO) https://github.com/apache/beam/issues/25799
    // uncomment tfx bsl tests once tfx supports protobuf 4.x
   // 'tfxInferenceTest',