Add testing workflow and use pytest

.github/workflows/ci-test.yml

@@ -0,0 +1,40 @@
+# Github action definitions for unit-tests with PRs.
+
+name: tft-unit-tests
+on:
+  pull_request:
+    branches: [ master ]
+    paths-ignore:
+      - '**.md'
+      - 'docs/**'
+  workflow_dispatch:
+
+jobs:
+  unit-tests:
+    if: github.actor != 'copybara-service[bot]'
+    runs-on: ubuntu-latest
+    timeout-minutes: 60
+
+    strategy:
+      matrix:
+        python-version: ['3.9', '3.10', '3.11']
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+        cache: 'pip'
+        cache-dependency-path: |
+          setup.py
+
+    - name: Install dependencies
+      run: |
+        pip install .[test]
+
+    - name: Run unit tests
+      shell: bash
+      run: |
+        pytest 

README.md

@@ -42,23 +42,28 @@ pip install tensorflow-transform
 To build from source follow the following steps:
 Create a virtual environment by running the commands
 
-```
-python3 -m venv <virtualenv_name>
+```bash
+python -m venv <virtualenv_name>
 source <virtualenv_name>/bin/activate
-pip3 install setuptools wheel
 git clone https://github.com/tensorflow/transform.git
 cd transform
-python3 setup.py bdist_wheel
+pip install .
 ```
 
-This will build the TFT wheel in the dist directory. To install the wheel from
-dist directory run the commands
+If you are doing development on the TFT repo, replace
 
+```bash
+pip install .
 ```
-cd dist
-pip3 install tensorflow_transform-<version>-py3-none-any.whl
+
+with
+
+```bash
+pip install -e .
 ```
 
+The `-e` flag causes TFT to be installed in [development mode](https://setuptools.pypa.io/en/latest/userguide/development_mode.html).
+
 ### Nightly Packages
 
 TFT also hosts nightly packages at https://pypi-nightly.tensorflow.org on
@@ -72,6 +77,14 @@ pip install --extra-index-url https://pypi-nightly.tensorflow.org/simple tensorf
 This will install the nightly packages for the major dependencies of TFT such
 as TensorFlow Metadata (TFMD), TFX Basic Shared Libraries (TFX-BSL).
 
+### Running Tests
+
+To run TFT tests, run the following command from the root of the repository:
+
+```bash
+pytest
+```
+
 ### Notable Dependencies
 
 TensorFlow is required.

pytest.ini

@@ -0,0 +1,5 @@
+[pytest]
+addopts = --verbose --import-mode=importlib
+testpaths = tensorflow_transform
+python_files = *_test.py
+norecursedirs = .* *.egg

setup.py

@@ -19,98 +19,100 @@
 
 
 def select_constraint(default, nightly=None, git_master=None):
-  """Select dependency constraint based on TFX_DEPENDENCY_SELECTOR env var."""
-  selector = os.environ.get('TFX_DEPENDENCY_SELECTOR')
-  if selector == 'UNCONSTRAINED':
-    return ''
-  elif selector == 'NIGHTLY' and nightly is not None:
-    return nightly
-  elif selector == 'GIT_MASTER' and git_master is not None:
-    return git_master
-  else:
-    return default
+    """Select dependency constraint based on TFX_DEPENDENCY_SELECTOR env var."""
+    selector = os.environ.get("TFX_DEPENDENCY_SELECTOR")
+    if selector == "UNCONSTRAINED":
+        return ""
+    elif selector == "NIGHTLY" and nightly is not None:
+        return nightly
+    elif selector == "GIT_MASTER" and git_master is not None:
+        return git_master
+    else:
+        return default
 
 
 # Get version from version module.
-with open('tensorflow_transform/version.py') as fp:
-  globals_dict = {}
-  exec(fp.read(), globals_dict)  # pylint: disable=exec-used
-__version__ = globals_dict['__version__']
+with open("tensorflow_transform/version.py") as fp:
+    globals_dict = {}
+    exec(fp.read(), globals_dict)  # pylint: disable=exec-used
+__version__ = globals_dict["__version__"]
 
 
 def _make_required_install_packages():
-  # Make sure to sync the versions of common dependencies (absl-py, numpy, and
-  # protobuf) with TF and pyarrow version with tfx-bsl.
-  return [
-      'absl-py>=0.9,<2.0.0',
-      'apache-beam[gcp]>=2.53,<3;python_version>="3.11"',
-      'apache-beam[gcp]>=2.47,<3;python_version<"3.11"',
-      'numpy>=1.22.0',
-      'protobuf>=4.25.2,<5;python_version>="3.11"',
-      'protobuf>=3.20.3,<5;python_version<"3.11"',
-      'pyarrow>=10,<11',
-      'pydot>=1.2,<2',
-      'tensorflow'
-      + select_constraint(
-          default='>=2.15,<2.16',
-          nightly='>=2.16.0.dev',
-          git_master='@git+https://github.com/tensorflow/tensorflow@master',
-      ),
-      'tensorflow-metadata'
-      + select_constraint(
-          default='>=1.15.0,<1.16.0',
-          nightly='>=1.16.0.dev',
-          git_master='@git+https://github.com/tensorflow/metadata@master',
-      ),
-      'tfx-bsl'
-      + select_constraint(
-          default='>=1.15.1,<1.16.0',
-          nightly='>=1.16.0.dev',
-          git_master='@git+https://github.com/tensorflow/tfx-bsl@master',
-      ),
-  ]
+    # Make sure to sync the versions of common dependencies (absl-py, numpy, and
+    # protobuf) with TF and pyarrow version with tfx-bsl.
+    return [
+        "absl-py>=0.9,<2.0.0",
+        'apache-beam[gcp]>=2.53,<3;python_version>="3.11"',
+        'apache-beam[gcp]>=2.47,<3;python_version<"3.11"',
+        "numpy>=1.22.0",
+        'protobuf>=4.25.2,<5;python_version>="3.11"',
+        'protobuf>=3.20.3,<5;python_version<"3.11"',
+        "pyarrow>=10,<11",
+        "pydot>=1.2,<2",
+        "tensorflow"
+        + select_constraint(
+            default=">=2.15,<2.16",
+            nightly=">=2.16.0.dev",
+            git_master="@git+https://github.com/tensorflow/tensorflow@master",
+        ),
+        "tensorflow-metadata"
+        + select_constraint(
+            default=">=1.15.0,<1.16.0",
+            nightly=">=1.16.0.dev",
+            git_master="@git+https://github.com/tensorflow/metadata@master",
+        ),
+        "tfx-bsl"
+        + select_constraint(
+            default=">=1.15.1,<1.16.0",
+            nightly=">=1.16.0.dev",
+            git_master="@git+https://github.com/tensorflow/tfx-bsl@master",
+        ),
+    ]
 
 
 # Get the long description from the README file.
-with open('README.md') as fp:
-  _LONG_DESCRIPTION = fp.read()
+with open("README.md") as fp:
+    _LONG_DESCRIPTION = fp.read()
 
 setup(
-    name='tensorflow-transform',
+    name="tensorflow-transform",
     version=__version__,
-    author='Google Inc.',
-    author_email='[email protected]',
-    license='Apache 2.0',
+    author="Google Inc.",
+    author_email="[email protected]",
+    license="Apache 2.0",
     classifiers=[
-        'Development Status :: 5 - Production/Stable',
-        'Intended Audience :: Developers',
-        'Intended Audience :: Education',
-        'Intended Audience :: Science/Research',
-        'License :: OSI Approved :: Apache Software License',
-        'Operating System :: OS Independent',
-        'Programming Language :: Python',
-        'Programming Language :: Python :: 3',
-        'Programming Language :: Python :: 3.9',
-        'Programming Language :: Python :: 3.10',
-        'Programming Language :: Python :: 3.11',
-        'Programming Language :: Python :: 3 :: Only',
-        'Topic :: Scientific/Engineering',
-        'Topic :: Scientific/Engineering :: Artificial Intelligence',
-        'Topic :: Scientific/Engineering :: Mathematics',
-        'Topic :: Software Development',
-        'Topic :: Software Development :: Libraries',
-        'Topic :: Software Development :: Libraries :: Python Modules',
+        "Development Status :: 5 - Production/Stable",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Education",
+        "Intended Audience :: Science/Research",
+        "License :: OSI Approved :: Apache Software License",
+        "Operating System :: OS Independent",
+        "Programming Language :: Python",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3 :: Only",
+        "Topic :: Scientific/Engineering",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Scientific/Engineering :: Mathematics",
+        "Topic :: Software Development",
+        "Topic :: Software Development :: Libraries",
+        "Topic :: Software Development :: Libraries :: Python Modules",
     ],
     namespace_packages=[],
     install_requires=_make_required_install_packages(),
-    python_requires='>=3.9,<4',
+    extras_require={"test": ["pytest>=8.0"]},
+    python_requires=">=3.9,<4",
     packages=find_packages(),
     include_package_data=True,
-    package_data={'tensorflow_transform': ['py.typed']},
-    description='A library for data preprocessing with TensorFlow',
+    package_data={"tensorflow_transform": ["py.typed"]},
+    description="A library for data preprocessing with TensorFlow",
     long_description=_LONG_DESCRIPTION,
-    long_description_content_type='text/markdown',
-    keywords='tensorflow transform tfx',
-    url='https://www.tensorflow.org/tfx/transform/get_started',
-    download_url='https://github.com/tensorflow/transform/tags',
-    requires=[])
+    long_description_content_type="text/markdown",
+    keywords="tensorflow transform tfx",
+    url="https://www.tensorflow.org/tfx/transform/get_started",
+    download_url="https://github.com/tensorflow/transform/tags",
+    requires=[],
+)

tensorflow_transform/analyzers_test.py

@@ -624,5 +624,3 @@ def testMinDiffFromAvg(self):
         analyzers.calculate_recommended_min_diff_from_avg(100000000), 25)
 
 
-if __name__ == '__main__':
-  test_case.main()

tensorflow_transform/annotators_test.py

@@ -65,5 +65,3 @@ def preprocessing_fn():
     self.assertEqual(trackable_object, object_tracker.trackable_objects[0])
 
 
-if __name__ == '__main__':
-  test_case.main()

tensorflow_transform/beam/analysis_graph_builder_test.py

@@ -579,5 +579,3 @@ class _Analyzer(
                                        structured_outputs)
 
 
-if __name__ == '__main__':
-  tft_unit.main()

tensorflow_transform/beam/analyzer_cache_test.py

@@ -310,5 +310,3 @@ def expand(self, pbegin):
             beam_test_util.equal_to([test_cache_dict[key].cache_dict['b']]))
 
 
-if __name__ == '__main__':
-  test_case.main()

tensorflow_transform/beam/analyzer_impls_test.py

@@ -136,5 +136,3 @@ def testJoinBoundarieRows(self, input_boundaries, expected_boundaries,
     self.assertAllEqual(num_buckets, expected_num_buckets)
 
 
-if __name__ == '__main__':
-  tft_unit.main()

tensorflow_transform/beam/annotators_test.py

@@ -259,5 +259,3 @@ def preprocessing_fn(inputs):
     )
 
 
-if __name__ == '__main__':
-  tft_unit.main()

tensorflow_transform/beam/bucketize_integration_test.py

@@ -888,5 +888,3 @@ def testBucketizationSpecificDistribution(self):
         inputs, expected_boundaries, tf.float32, num_buckets=5)
 
 
-if __name__ == '__main__':
-  tft_unit.main()

tensorflow_transform/beam/cached_impl_test.py

@@ -1749,5 +1749,3 @@ def preprocessing_fn(inputs):
     )
 
 
-if __name__ == '__main__':
-  tft_unit.main()

tensorflow_transform/beam/combiner_packing_util_test.py

@@ -433,5 +433,3 @@ def _side_effect_fn(saved_model_future, cache_value_nodes,
         second=expected_dot_graph_str_after_packing)
 
 
-if __name__ == '__main__':
-  test_case.main()

tensorflow_transform/beam/context_test.py

@@ -40,5 +40,3 @@ def testNestedContextCreateBaseTempDir(self):
       tft_beam.Context.create_base_temp_dir()
 
 
-if __name__ == '__main__':
-  tft_unit.main()

tensorflow_transform/beam/deep_copy_test.py

@@ -330,5 +330,3 @@ def testDeepCopyTags(self):
     self.assertEqual(DeepCopyTest._counts['Add2'], 3 * (num_copies + 1))
     self.assertEqual(DeepCopyTest._counts['Add3'], 3)
 
-if __name__ == '__main__':
-  unittest.main()

tensorflow_transform/beam/impl_output_record_batches_test.py

@@ -199,5 +199,3 @@ def testConvertToLargeRecordBatch(
     self.assertGreater(actual_num_batches, 1)
 
 
-if __name__ == '__main__':
-  tft_unit.main()

tensorflow_transform/beam/impl_test.py

@@ -4797,5 +4797,3 @@ def test_preprocessing_fn_returns_wrong_type(self):
           expected_data=None)
 
 
-if __name__ == '__main__':
-  tft_unit.main()

tensorflow_transform/beam/tft_beam_io/beam_metadata_io_test.py

@@ -92,5 +92,3 @@ def mock_write_metadata(metadata, path):
       self.assertEqual(metadata, test_metadata.COMPLETE_METADATA)
 
 
-if __name__ == '__main__':
-  tf.test.main()

tensorflow_transform/beam/tft_beam_io/transform_fn_io_test.py

@@ -149,5 +149,3 @@ def mock_copy_tree_to_unique_temp_dir(source, base_temp_dir_path):
     self.assertEqual(2, len(file_io.list_directory(transform_output_dir)))
 
 
-if __name__ == '__main__':
-  tf.test.main()

tensorflow_transform/beam/tukey_hh_params_integration_test.py

@@ -627,5 +627,3 @@ def assert_and_cast_dtype(tensor):
         # Runs the test deterministically on the whole batch.
         beam_pipeline=beam.Pipeline())
 
-if __name__ == '__main__':
-  tft_unit.main()

tensorflow_transform/beam/vocabulary_integration_test.py

@@ -2088,5 +2088,3 @@ def preprocessing_fn(inputs):
     )
 
 
-if __name__ == '__main__':
-  tft_unit.main()

tensorflow_transform/beam/vocabulary_tfrecord_gzip_integration_test.py

@@ -25,5 +25,3 @@ def _VocabFormat(self):
     return 'tfrecord_gzip'
 
 
-if __name__ == '__main__':
-  tft_unit.main()

tensorflow_transform/coders/csv_coder_test.py

@@ -292,5 +292,3 @@ def test_picklable(self):
       self.assertEqual(coder.encode(instance), csv_line.encode('utf-8'))
 
 
-if __name__ == '__main__':
-  test_case.main()

tensorflow_transform/coders/example_proto_coder_test.py

@@ -406,5 +406,3 @@ def test_example_proto_coder_cache(self):
       self.assertSerializedProtosEqual(coder.encode(instance), serialized_proto)
 
 
-if __name__ == '__main__':
-  test_case.main()