sdks/python/tox.ini

;
;    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.
;

[tox]
# new environments will be excluded by default unless explicitly added to envlist.
envlist = py311,py312,py313,py314,py311-{cloud,cloudcoverage,dask},py312-{cloud,dask},py313-{cloud,dask},py314-{cloud,dask},docs,lint,whitespacelint
toxworkdir = {toxinidir}/target/{env:ENV_NAME:.tox}

[pycodestyle]
# Disable all errors and warnings except for the ones related to blank lines.
# pylint does not check the number of blank lines.
select = E3

# Shared environment options.
[testenv]
# Run the tests using pre-released dependencies.
# https://github.com/apache/beam/issues/25668
pip_pre = True
# allow apps that support color to use it.
passenv=TERM,CLOUDSDK_CONFIG,DOCKER_*,TESTCONTAINERS_*,TC_*,ALLOYDB_PASSWORD
# Set [] options for pip installation of apache-beam tarball.
extras = test,dataframe,hadoop,redis,tfrecord,yaml
# Don't warn that these commands aren't installed.
allowlist_externals =
  false
  time
  bash
  rm
  /bin/sh
  curl
  ./codecov
  chmod
setenv =
  RUN_SKIPPED_PY3_TESTS=0
  # Set before importing gRPC; test-class setup is too late.
  # TODO(https://github.com/grpc/grpc/issues/37710): Remove once fixed.
  GRPC_ENABLE_FORK_SUPPORT=0
  # Use an isolated tmp dir for tests that get slowed down by scanning /tmp.
  TMPDIR={envtmpdir}
  # Silence warning about ignoring PYTHONPATH.
  PYTHONPATH=

# These 2 magic command overrides are required for Jenkins builds.
# Otherwise we get "OSError: [Errno 2] No such file or directory" errors.
# Source:
# https://github.com/tox-dev/tox/issues/123#issuecomment-284714629
install_command = {envbindir}/python {envbindir}/pip install --retries 10 {opts} {packages}
list_dependencies_command = {envbindir}/python {envbindir}/pip freeze
commands_pre =
  python --version
  pip --version
  pip check
  bash {toxinidir}/scripts/run_tox_cleanup.sh
commands_post =
  bash {toxinidir}/scripts/run_tox_cleanup.sh

commands = false {envname} is misconfigured

[testenv:py{311,312,313,314}]
commands_pre =
  python --version
  pip --version
  pip check
  bash {toxinidir}/scripts/run_tox_cleanup.sh
deps =
  numpy==1.26.4
commands =
  python apache_beam/examples/complete/autocomplete_test.py
  bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"

[testenv:py{311,312,313,314}-macos]
commands_pre =
  python --version
  pip --version
  # pip check
  bash {toxinidir}/scripts/run_tox_cleanup.sh
commands =
  python apache_beam/examples/complete/autocomplete_test.py
  bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"

[testenv:py{311,312,313,314}-win]
commands =
  python apache_beam/examples/complete/autocomplete_test.py
  bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"
install_command = {envbindir}/python.exe {envbindir}/pip.exe install --retries 10 {opts} {packages}
list_dependencies_command = {envbindir}/python.exe {envbindir}/pip.exe freeze

[testenv:py{311,312,313,314}-cloud]
; extras = test,gcp,interactive,dataframe,aws,azure
extras = test,hadoop,gcp,interactive,dataframe,aws,azure
commands =
  python apache_beam/examples/complete/autocomplete_test.py
  bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"

[testenv:py311-ml]
# Don't set TMPDIR to avoid "AF_UNIX path too long" errors in certain tests.
deps =
  pip==25.0.1
  accelerate>=1.6.0
  onnx<2.0.0
  onnxscript>=0.5.0,<1
setenv =
extras = test,gcp,dataframe,p311_ml_test
commands =
  # Log tensorflow version for debugging
  /bin/sh -c "pip freeze | grep -E tensorflow"
  bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"

[testenv:py312-ml]
# many packages do not support py3.12
# Don't set TMPDIR to avoid "AF_UNIX path too long" errors in certain tests.
deps =
  accelerate>=1.6.0
setenv =
extras = test,gcp,dataframe,p312_ml_test
commands =
  # Log tensorflow version for debugging
  /bin/sh -c "pip freeze | grep -E tensorflow"
  bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"

[testenv:py313-ml]
# many packages do not support py3.13, and datatables breaks after 3.12.
# Don't set TMPDIR to avoid "AF_UNIX path too long" errors in certain tests.
deps =
  accelerate>=1.6.0
setenv =
extras = test,gcp,dataframe,p313_ml_test
commands =
  # Log tensorflow version for debugging
  /bin/sh -c "pip freeze | grep -E tensorflow"
  bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"

[testenv:py314-ml]
deps =
  accelerate>=1.6.0
setenv =
extras = test,gcp,dataframe,p314_ml_test
commands =
  # Log tensorflow version for debugging
  /bin/sh -c "pip freeze | grep -E tensorflow"
  bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"

[testenv:py{311,312,313,314}-dask]
extras = test,dask,dataframe
commands_pre =
  pip install 'distributed>=2024.4.2' 'dask>=2024.4.2'
commands =
  bash {toxinidir}/scripts/run_pytest.sh {envname} {toxinidir}/apache_beam/runners/dask/

[testenv:py{311,312,313,314}-win-dask]
# use the tight range since the latest dask requires cloudpickle 3.0
commands_pre =
  pip install 'distributed>=2024.4.2,<2024.9.0' 'dask>=2024.4.2,<2024.9.0'
commands =
  python apache_beam/examples/complete/autocomplete_test.py
  bash {toxinidir}/scripts/run_pytest.sh {envname} {toxinidir}/apache_beam/runners/dask/
install_command = {envbindir}/python.exe {envbindir}/pip.exe install --retries 10 {opts} {packages}
list_dependencies_command = {envbindir}/python.exe {envbindir}/pip.exe freeze

[testenv:py311-cloudcoverage]
deps =
  pytest-cov==3.0.0

platform = linux
passenv = GIT_*,BUILD_*,ghprb*,CHANGE_ID,BRANCH_NAME,JENKINS_*,CODECOV_*,GITHUB_*,DOCKER_*,TESTCONTAINERS_*,TC_*

# Don't set TMPDIR to avoid "AF_UNIX path too long" errors in certain tests.
setenv =
  PYTHONPATH = {toxinidir}
  DOCKER_HOST = {env:DOCKER_HOST}
  TC_TIMEOUT = {env:TC_TIMEOUT:120}
  TC_MAX_TRIES = {env:TC_MAX_TRIES:120}
  TC_SLEEP_TIME = {env:TC_SLEEP_TIME:1}

# NOTE: we could add ml_test to increase the collected code coverage metrics, but it would make the suite slower.
extras = test,hadoop,gcp,interactive,dataframe,aws,redis
commands =
  bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}" "--cov-report=xml --cov=. --cov-append"

[testenv:lint]
# Don't set TMPDIR to avoid "AF_UNIX path too long" errors in pylint.
setenv =
deps =
  dask==2022.01.0
  distributed==2022.01.0
extras =
  gcp
  dev
commands =
  ruff --version
  time {toxinidir}/scripts/run_lint.sh
  pyrefly --version
  time python setup.py pyrefly

[testenv:whitespacelint]
setenv =
deps =
  whitespacelint==1.1.0
commands =
  time {toxinidir}/scripts/run_whitespacelint.sh

[testenv:docs]
extras = test,gcp,docs,interactive,dataframe,dask
deps =
  Sphinx==7.4.7
  sphinx_rtd_theme==3.0.1
  docutils>=0.18.1
  Jinja2==3.1.0
commands =
  time {toxinidir}/scripts/generate_pydoc.sh

[testenv:hdfs_integration_test]
# Used by hdfs_integration_test.sh. Do not run this directly, as it depends on
# nodes defined in hdfs_integration_test/docker-compose.yml.
deps =
  holdup==1.8.0
extras =
  gcp
  hadoop
allowlist_externals =
  bash
  echo
  sleep
  wget
  hdfscli
passenv = HDFSCLI_CONFIG
commands =
  holdup -t 45 http://namenode:50070 http://datanode:50075
  echo "Waiting for safe mode to end."
  sleep 45
  wget storage.googleapis.com/dataflow-samples/shakespeare/kinglear.txt
  hdfscli -v -v -v upload -f kinglear.txt /
  python -m apache_beam.examples.wordcount \
      --input hdfs://kinglear* \
      --output hdfs://py-wordcount-integration \
      --hdfs_host namenode --hdfs_port 50070 --hdfs_user root
  python -m apache_beam.examples.wordcount \
      --input hdfs://unused_server/kinglear* \
      --output hdfs://unused_server/py-wordcount-integration \
      --hdfs_host namenode --hdfs_port 50070 --hdfs_user root --hdfs_full_urls
commands_pre =
  pip check

[testenv:azure_integration_test]
# Used by azure/integration_test/azure_integration_test.sh.
# Do not run this directly, as it depends on nodes defined in
# azure/integration_test/docker-compose.yml.
deps =
extras =
  azure
passenv = REQUESTS_CA_BUNDLE
allowlist_externals =
  wget
  az
  bash
setenv =
  CONNECTION_STRING="DefaultEndpointsProtocol=https;AccountName=devstoreaccount1;AccountKey=Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==;BlobEndpoint=https://azurite:10000/devstoreaccount1;"
commands_pre =
  pip check
  wget storage.googleapis.com/dataflow-samples/shakespeare/kinglear.txt
  # Create container for storing files.
  az storage container create -n container --connection-string {env:CONNECTION_STRING}
  # Upload test file.
  az storage blob upload -f kinglear.txt -c container -n kinglear.txt --connection-string {env:CONNECTION_STRING}
commands =
  # Test --azure_connection_string
  python -m apache_beam.examples.wordcount \
      --input azfs://devstoreaccount1/container/* \
      --output azfs://devstoreaccount1/container/py-wordcount-integration \
      --azure_connection_string {env:CONNECTION_STRING}
  # This doesn't work because there's no way to send a fake bearer token to
  # Azurite when using DefaultAzureCredential.
  # See https://github.com/Azure/Azurite/issues/389#issuecomment-615298432
  # and https://github.com/Azure/Azurite/issues/1750#issue-1449778593
  #python -m apache_beam.examples.wordcount \
  #    --input azfs://devstoreaccount1/container/* \
  #    --output azfs://devstoreaccount1/container/py-wordcount-integration \
  #    --blob_service_endpoint https://azurite:10000/devstoreaccount1/container-name \
  #    --azure_managed_identity_client_id "abc123"

[testenv:py3-yapf]
# keep the version of yapf in sync with the 'rev' in .pre-commit-config.yaml and pyproject.toml
deps =
  yapf==0.43.0
commands =
  yapf --version
  time yapf --in-place --parallel --recursive apache_beam

[testenv:py3-yapf-check]
# keep the version of yapf in sync with the 'rev' in .pre-commit-config.yaml and pyproject.toml
deps =
  yapf==0.43.0
commands =
  yapf --version
  time yapf --diff --parallel --recursive apache_beam

[testenv:jest]
setenv =
deps =
  jupyterlab==4.3.6
commands =
  time {toxinidir}/scripts/setup_nodejs.sh
  time {toxinidir}/scripts/run_jest.sh

[testenv:eslint]
setenv =
deps =
  jupyterlab==4.3.6
commands =
  time {toxinidir}/scripts/setup_nodejs.sh
  time {toxinidir}/scripts/run_eslint.sh

[testenv:flink-runner-test]
extras = test
commands =
  bash {toxinidir}/scripts/pytest_validates_runner.sh {envname} {toxinidir}/apache_beam/runners/portability/flink_runner_test.py {posargs}

[testenv:spark-runner-test]
extras = test
commands =
  bash {toxinidir}/scripts/pytest_validates_runner.sh {envname} {toxinidir}/apache_beam/runners/portability/spark_runner_test.py {posargs}

[testenv:kafka-streams-runner-test]
extras = test
commands =
  bash {toxinidir}/scripts/pytest_validates_runner.sh {envname} {toxinidir}/apache_beam/runners/portability/kafka_streams_runner_test.py {posargs}

[testenv:prism-runner-test]
extras = test
commands =
  bash {toxinidir}/scripts/pytest_validates_runner.sh {envname} {toxinidir}/apache_beam/runners/portability/prism_runner_test.py {posargs}

[testenv:py311-pyarrow-{17,18,19,20,21,22,23,24,25}]
deps =
  # Test against versions of pyarrow released in last ~2 years.
  17: pyarrow>=17,<18
  18: pyarrow>=18,<19
  19: pyarrow>=19,<20
  20: pyarrow>=20,<21
  21: pyarrow>=21,<22
  22: pyarrow>=22,<23
  23: pyarrow>=23,<24
  24: pyarrow>=24,<25
  25: pyarrow>=25,<26
  numpy==1.26.4
commands =
  # Log pyarrow and numpy version for debugging
  /bin/sh -c "pip freeze | grep -E '(pyarrow|numpy)'"
  # Run pytest directly rather using run_pytest.sh. It doesn't handle
  # selecting tests with -m (BEAM-12985).
  # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
  /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_pyarrow {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'


[testenv:py311-pandas-{22,23}]
# Keep in line with dataframe_dependency in setup.py. pandas 2.2.2 is the first
# release compatible with numpy 2.
deps =
  # Lower bound of the supported range; also keeps numpy 1.x covered.
  22: pandas>=2.2.2,<2.3.0
  22: numpy>=1.23.2,<2.0.0
  23: pandas>=2.3.0,<2.4.0
commands =
  # Log pandas and numpy version for debugging
  /bin/sh -c "pip freeze | grep -E '(pandas|numpy)'"
  # Run all DataFrame API unit tests
  bash {toxinidir}/scripts/run_pytest.sh {envname} 'apache_beam/dataframe'

[testenv:py311-tft-121]
# Keep in line with the tft extra in setup.py. tensorflow-transform<1.21
# requires pyarrow<11, which is incompatible with Beam's pyarrow>=14.
deps =
  tensorflow_transform>=1.21.0,<1.22.0
  dill
extras = test,gcp
commands =
  bash {toxinidir}/scripts/run_pytest.sh {envname} 'apache_beam/ml/transforms apache_beam/examples/snippets/transforms/elementwise/mltransform_test.py'

[testenv:py311-pytorch-{23,214}]
# Keep in line with the torch bounds in setup.py. torch<2.3 is not compatible
# with numpy 2, so numpy is intentionally left unpinned.
deps =
  23: torch>=2.3.0,<2.4.0
  214: torch>=2.14.0,<2.15.0
extras = test,gcp
# Don't set TMPDIR to avoid "AF_UNIX path too long" errors in certain tests.
setenv =
commands =
  # Log torch version for debugging
  /bin/sh -c "pip freeze | grep -E torch"
  # Run all PyTorch unit tests
  # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
  /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_pytorch {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'

[testenv:py314-pytorch-214]
# Keep in line with py311-pytorch-latest above. torch>=2.10 is required for
# Python 3.14, so the torch 2.3 lower-bound env only runs on Python 3.11.
deps =
  torch>=2.14.0,<2.15.0
extras = test,gcp
# Don't set TMPDIR to avoid "AF_UNIX path too long" errors in certain tests.
setenv =
commands =
  # Log torch version for debugging
  /bin/sh -c "pip freeze | grep -E torch"
  # Run all PyTorch unit tests
  # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
  /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_pytorch {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'

[testenv:py311-onnx-{119,latest}]
# Keep in line with the onnx extra in setup.py. tensorflow is pinned to 2.21 in
# both environments because tensorflow<2.20 cannot be installed alongside the
# gcp extra (protobuf<6 vs protobuf>=6.33.5).
deps =
  # Lower bound of the supported range.
  119: onnxruntime>=1.19.0,<1.20
  119: onnx>=1.16.0,<1.17
  119: skl2onnx>=1.17.0,<1.18
  119: torch>=2.3.0,<2.4.0
  latest: onnxruntime>=1.31.0,<2
  latest: onnx>=1.23.0,<2
  latest: skl2onnx>=1.20.0,<2
  latest: torch>=2.11.0,<2.12.0
  tf2onnx>=1.17.0,<1.18
  tensorflow>=2.21.0,<2.22
  # torch>=2.9 exports to ONNX through torch.export, which requires onnxscript.
  onnxscript>=0.5.0,<1
extras = test,gcp
commands =
  # Log onnx version for debugging
  /bin/sh -c "pip freeze | grep -E onnx"
  # Run all ONNX unit tests
  # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
  /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_onnx {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'

[testenv:py311-tensorflow-{216,221}]
# Keep in line with the tensorflow extra in setup.py.
# tensorflow<2.20 requires protobuf<6, which cannot be installed alongside the
# gcp extra (protobuf>=6.33.5 via google-cloud-bigtable), so the lower-bound
# environment only installs the test extra.
# tensorflow-hub is needed by the TensorFlow inference tests; it imports
# pkg_resources, which was removed in setuptools 82.
deps =
  216: tensorflow>=2.16.0,<2.17
  216: tf-keras>=2.16.0,<2.17
  216: protobuf>=4.25.2,<5
  221: tensorflow>=2.21.0,<2.22
  221: tf-keras>=2.21.0,<2.22
  tensorflow-hub>=0.16.1,<0.17
  setuptools<82
extras =
  216: test
  221: test,gcp
# Apply the deps above as constraints when installing Beam's own dependencies,
# so that they cannot upgrade e.g. protobuf past tensorflow 2.16's protobuf<5
# cap.
constrain_package_deps = true
commands =
  # Log tensorflow version for debugging
  /bin/sh -c "pip freeze | grep -E tensorflow"
  # Run all Tensorflow unit tests
  # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
  /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_tf {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'

[testenv:py311-xgboost-{200,300}]
# Keep these bounds in line with xgboost_dependency in setup.py.
deps =
  200: xgboost>=2.0.0,<3.0.0
  300: xgboost>=3.0.0,<4
  datatable>=1.1.0,<2
extras = test,gcp
commands =
  # Log XGBoost version for debugging
  /bin/sh -c "pip freeze | grep -E xgboost"
  # Run all XGBoost unit tests
  # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
  /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_xgboost {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'

[testenv:py311-transformers-{441,457,latest}]
deps =
  # Environment dependencies are defined in the `setenv` section and installed in the `commands` section.
pip_pre = False
extras = test
setenv =
    # Keep in line with the transformers extra in setup.py.
    # tensorflow-hub is needed by the TensorFlow handler tests; it imports
    # pkg_resources, which was removed in setuptools 82.
    COMMON_DEPS = 'tensorflow>=2.16,<2.22' tf-keras tensorflow-hub 'setuptools<82' sentence-transformers pip==25.0.1
    # Lower bound of the supported range.
    441: DEPS = 'transformers>=4.41.0,<4.42.0' 'torch>=2.3.0,<2.4.0'
    # Last transformers release line with TensorFlow support.
    # torch<2.12: torch>=2.12 pulls in triton>=3.7, which crashes when imported
    # after tensorflow.
    457: DEPS = 'transformers>=4.57.0,<4.58.0' 'torch>=2.3.0,<2.12.0'
    latest: DEPS = 'transformers>=5.20.0,<5.21.0' 'torch>=2.11.0,<2.12.0' 'accelerate>=1.6.0'
commands =
  /bin/sh -c "pip install .[{extras}] {env:DEPS} {env:COMMON_DEPS}"
  # Log transformers and its dependencies version for debugging
  /bin/sh -c "pip freeze | grep -E transformers"
  /bin/sh -c "pip freeze | grep -E torch"
  /bin/sh -c "pip freeze | grep -E tensorflow"
  # Run all Transformers unit tests
  # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
  /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_transformers {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'

[testenv:py{311,313}-vertex-ai]
deps =
  tensorflow==2.12.0
extras = test,gcp
commands =
  # Log aiplatform and its dependencies version for debugging
  /bin/sh -c "pip freeze | grep -E google-cloud-aiplatform"
  /bin/sh -c "pip freeze | grep -E tensorflow"
  # Run all Vertex AI unit tests
  # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
  /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_vertex_ai {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'


[testenv:py311-embeddings]
deps =
  sentence-transformers==3.3.1
  accelerate>=1.6.0
  transformers>=4.28.0,<4.56.0
passenv = HF_INFERENCE_TOKEN
extras = test,gcp
commands =
  # Log aiplatform and its dependencies version for debugging
  /bin/sh -c "pip freeze | grep -E sentence-transformers"
  /bin/sh -c "pip freeze | grep -E google-cloud-aiplatform"
  /bin/sh -c "pip freeze | grep -E transformers"
  # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
  /bin/sh -c 'pytest apache_beam/ml/transforms/embeddings -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'


[testenv:py311-TFHubEmbeddings-016]
# Keep in line with the tensorflow-hub extra in setup.py. tensorflow-hub<0.16
# imports tf.compat.v1.estimator, which was removed in tensorflow 2.16.
deps =
  tensorflow-hub>=0.16.1,<0.17
  tensorflow-text # required to register ops for text embedding models.
  # tensorflow-hub imports pkg_resources, which was removed in setuptools 82.
  setuptools<82

extras = test,gcp
commands =
  # Log aiplatform and its dependencies version for debugging
  /bin/sh -c "pip freeze | grep -E tensorflow"
  # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
  bash {toxinidir}/scripts/run_pytest.sh {envname} 'apache_beam/ml/transforms/embeddings'

[testenv:py{311,312}-dill]
extras = test,dill
commands =
  # Log dill version for debugging
  /bin/sh -c "pip freeze | grep -E dill"
  # Run all dill-specific tests
  bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"