Re enable more not e2e xfail test cases (#6947)

* Exclude tests for the deprecated tfx/orchestration/experimental/core module from experimental orchestration.
* Include additional testdata files in the package build.
* Add module-level pytest cleanup functions for proto classes.
* Update doc tests to clean up docs before execution.

Author: nikelite

MANIFEST.in

@@ -14,5 +14,8 @@ include tfx/proto/*.proto
 recursive-include tfx/orchestration/kubeflow/v2/testdata *
 
 recursive-include tfx/components/testdata *
+recursive-include tfx/orchestration/kubeflow/v2/testdata *
 
-include tfx/examples/imdb/data/
+include tfx/examples/imdb/data/*
+include tfx/orchestration/beam/testdata/*
+include tfx/orchestration/kubeflow/v2/container/testdata/*

pyproject.toml

@@ -40,7 +40,7 @@ Repository = "https://github.com/tensorflow/tfx"
 addopts = "--import-mode=importlib"
 testpaths = "tfx"
 python_files = "*_test.py"
-norecursedirs = ["custom_components", ".*", "*.egg"]
+norecursedirs = ["custom_components", ".*", "*.egg", "tfx/orchestration/experimental/core"]
 markers = [
     "e2e: end-to-end tests which are slow and require more dependencies (deselect with '-m \"not end_to_end\"')",
     "serial: mark tests that should not run in parallel",

tfx/dsl/placeholder/proto_placeholder_test.py

@@ -15,7 +15,9 @@
 
 import base64
 import functools
+import importlib
 import os
+import pytest
 from typing import Any, Optional, TypeVar, Union
 
 import tensorflow as tf
@@ -34,6 +36,13 @@
 from google.protobuf import text_format
 from ml_metadata.proto import metadata_store_pb2
 
+
+
+@pytest.fixture(autouse=True,scope="module")
+def cleanup():
+    yield
+    importlib.reload(pipeline_pb2)
+
 _ExecutionInvocation = functools.partial(
     ph.make_proto, execution_invocation_pb2.ExecutionInvocation()
 )

tfx/examples/penguin/experimental/sklearn_predict_extractor_test.py

@@ -69,8 +69,8 @@ def setUp(self):
         self._makeExample(age=5.0, language=0.0, label=0),
     ]
 
-  @pytest.mark.xfail(run=False, reason="PR 6889 This test fails and needs to be fixed. "
-"If this test passes, please remove this mark.", strict=True)
+  @pytest.mark.xfail(run=False, reason="This is based on experimental implementation,"
+"and the test fails.", strict=True)
   def testMakeSklearnPredictExtractor(self):
     """Tests that predictions are made from extracts for a single model."""
     feature_extractor = tfma.extractors.FeaturesExtractor(self._eval_config)
@@ -98,8 +98,8 @@ def check_result(actual):
 
       util.assert_that(predict_extracts, check_result)
 
-  @pytest.mark.xfail(run=False, reason="PR 6889 This test fails and needs to be fixed. "
-"If this test passes, please remove this mark.", strict=True)
+  @pytest.mark.xfail(run=False, reason="This is based on experimental implementation,"
+"and the test fails.", strict=True)
   def testMakeSklearnPredictExtractorWithMultiModels(self):
     """Tests that predictions are made from extracts for multiple models."""
     eval_config = tfma.EvalConfig(model_specs=[

tfx/examples/ranking/struct2tensor_parsing_utils_test.py

@@ -15,7 +15,6 @@
 
 
 
-import pytest
 import itertools
 import unittest
 
@@ -172,15 +171,11 @@
 ]
 
 
-@pytest.mark.xfail(run=False, reason="PR 6889 This class contains tests that fail and needs to be fixed. "
-"If all tests pass, please remove this mark.")
 @unittest.skipIf(struct2tensor_parsing_utils is None,
                  'Cannot import required modules. This can happen when'
                  ' struct2tensor is not available.')
 class ELWCDecoderTest(tf.test.TestCase):
 
-  #@pytest.mark.xfail(run=False, reason="PR 6889 This test fails and needs to be fixed. "
-#"If this test passes, please remove this mark.", strict=True)
   def testAllDTypes(self):
     context_features = [
         struct2tensor_parsing_utils.Feature('ctx.int', tf.int64),

tfx/orchestration/beam/beam_dag_runner_test.py

@@ -14,7 +14,6 @@
 """Tests for tfx.orchestration.portable.beam_dag_runner."""
 
 
-import pytest
 import os
 from typing import Optional
 
@@ -172,8 +171,6 @@ def _run_node(self):
     _executed_components.append(self._node_id)
 
 
-@pytest.mark.xfail(run=False, reason="PR 6889 This class contains tests that fail and needs to be fixed. "
-"If all tests pass, please remove this mark.")
 class BeamDagRunnerTest(test_case_utils.TfxTest):
 
   def setUp(self):

tfx/orchestration/data_types_utils_test.py

@@ -14,6 +14,8 @@
 """Tests for tfx.orchestration.data_types_utils."""
 
 
+import importlib
+import pytest
 from absl.testing import parameterized
 from tfx import types
 from tfx.orchestration import data_types_utils
@@ -32,6 +34,12 @@
 _DEFAULT_ARTIFACT_TYPE_NAME = 'Examples'
 
 
+@pytest.fixture(scope="module", autouse=True)
+def cleanup():
+  yield
+  importlib.reload(struct_pb2)
+
+
 def _create_artifact(uri: str) -> types.Artifact:
   artifact = types.Artifact(
       metadata_store_pb2.ArtifactType(name=_DEFAULT_ARTIFACT_TYPE_NAME))

tfx/orchestration/kubeflow/v2/container/kubeflow_v2_entrypoint_utils_test.py

@@ -15,7 +15,6 @@
 
 
 
-import pytest
 import os
 from kfp.pipeline_spec import pipeline_spec_pb2 as pipeline_pb2
 import tensorflow as tf
@@ -68,8 +67,6 @@
 }
 
 
-@pytest.mark.xfail(run=False, reason="PR 6889 This class contains tests that fail and needs to be fixed. "
-"If all tests pass, please remove this mark.")
 class KubeflowV2EntrypointUtilsTest(tf.test.TestCase):
 
   def setUp(self):

tfx/orchestration/kubeflow/v2/container/kubeflow_v2_run_executor_test.py

@@ -14,7 +14,6 @@
 """Tests for kubeflow_v2_run_executor.py."""
 
 
-import pytest
 import json
 import os
 from typing import Any, Mapping, Sequence
@@ -100,8 +99,6 @@ def Do(self, input_dict: Mapping[str, Sequence[artifact.Artifact]],
 _EXEC_PROPERTIES = {"key_1": "value_1", "key_2": 536870911}
 
 
-@pytest.mark.xfail(run=False, reason="PR 6889 This class contains tests that fail and needs to be fixed. "
-"If all tests pass, please remove this mark.")
 class KubeflowV2RunExecutorTest(
     test_case_utils.TfxTest, parameterized.TestCase
 ):

tfx/orchestration/kubeflow/v2/file_based_example_gen/driver_test.py

@@ -14,7 +14,6 @@
 
 
 
-import pytest
 import json
 import os
 
@@ -93,8 +92,6 @@ def _load_test_file(filename: str):
   ).read()
 
 
-@pytest.mark.xfail(run=False, reason="PR 6889 This class contains tests that fail and needs to be fixed. "
-"If all tests pass, please remove this mark.")
 class RunDriverTest(test_case_utils.TfxTest, parameterized.TestCase):
 
   def setUp(self):