From 9e66ed10dc77f005568bfaf4aeb794109f33acb9 Mon Sep 17 00:00:00 2001
From: Doojin Park <doojin@google.com>
Date: Fri, 1 Nov 2024 13:52:46 +0900
Subject: [PATCH] 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.
---
 MANIFEST.in                                   |  5 +++-
 pyproject.toml                                |  2 +-
 tfx/dsl/placeholder/proto_placeholder_test.py |  9 +++++++
 .../sklearn_predict_extractor_test.py         |  8 +++---
 .../struct2tensor_parsing_utils_test.py       |  5 ----
 .../beam/beam_dag_runner_test.py              |  3 ---
 tfx/orchestration/data_types_utils_test.py    |  8 ++++++
 .../kubeflow_v2_entrypoint_utils_test.py      |  3 ---
 .../kubeflow_v2_run_executor_test.py          |  3 ---
 .../v2/file_based_example_gen/driver_test.py  |  3 ---
 .../local/local_dag_runner_test.py            |  9 -------
 .../local/local_pipeline_test.py              |  9 -------
 .../portable/inputs_utils_test.py             |  7 ------
 tfx/orchestration/portable/launcher_test.py   | 19 --------------
 .../portable/outputs_utils_test.py            | 11 --------
 .../portable/partial_run_utils_test.py        | 25 -------------------
 .../python_execution_binary_utils_test.py     |  3 ---
 tfx/types/artifact_test.py                    |  8 ++++++
 tfx/utils/doc_controls_test.py                |  5 ++--
 19 files changed, 36 insertions(+), 109 deletions(-)

diff --git a/MANIFEST.in b/MANIFEST.in
index f551d1da9b..d787ade117 100644
--- a/MANIFEST.in
+++ b/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/*
diff --git a/pyproject.toml b/pyproject.toml
index 9bf35696e2..10a6c6121d 100644
--- a/pyproject.toml
+++ b/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",
diff --git a/tfx/dsl/placeholder/proto_placeholder_test.py b/tfx/dsl/placeholder/proto_placeholder_test.py
index 2e03ec4f01..e36dce45f6 100644
--- a/tfx/dsl/placeholder/proto_placeholder_test.py
+++ b/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()
 )
diff --git a/tfx/examples/penguin/experimental/sklearn_predict_extractor_test.py b/tfx/examples/penguin/experimental/sklearn_predict_extractor_test.py
index 4b3e80c605..7fda470dc5 100644
--- a/tfx/examples/penguin/experimental/sklearn_predict_extractor_test.py
+++ b/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=[
diff --git a/tfx/examples/ranking/struct2tensor_parsing_utils_test.py b/tfx/examples/ranking/struct2tensor_parsing_utils_test.py
index 1785718e0d..f523ef1de7 100644
--- a/tfx/examples/ranking/struct2tensor_parsing_utils_test.py
+++ b/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),
diff --git a/tfx/orchestration/beam/beam_dag_runner_test.py b/tfx/orchestration/beam/beam_dag_runner_test.py
index 54bde196f0..79f06d3c26 100644
--- a/tfx/orchestration/beam/beam_dag_runner_test.py
+++ b/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):
diff --git a/tfx/orchestration/data_types_utils_test.py b/tfx/orchestration/data_types_utils_test.py
index 7b353054b6..41a842fed5 100644
--- a/tfx/orchestration/data_types_utils_test.py
+++ b/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))
diff --git a/tfx/orchestration/kubeflow/v2/container/kubeflow_v2_entrypoint_utils_test.py b/tfx/orchestration/kubeflow/v2/container/kubeflow_v2_entrypoint_utils_test.py
index 33b15b1777..ac8f0dc71f 100644
--- a/tfx/orchestration/kubeflow/v2/container/kubeflow_v2_entrypoint_utils_test.py
+++ b/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):
diff --git a/tfx/orchestration/kubeflow/v2/container/kubeflow_v2_run_executor_test.py b/tfx/orchestration/kubeflow/v2/container/kubeflow_v2_run_executor_test.py
index c956e0face..891f787b4b 100644
--- a/tfx/orchestration/kubeflow/v2/container/kubeflow_v2_run_executor_test.py
+++ b/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
 ):
diff --git a/tfx/orchestration/kubeflow/v2/file_based_example_gen/driver_test.py b/tfx/orchestration/kubeflow/v2/file_based_example_gen/driver_test.py
index 3900fb0af4..2d197d6e40 100644
--- a/tfx/orchestration/kubeflow/v2/file_based_example_gen/driver_test.py
+++ b/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):
diff --git a/tfx/orchestration/local/local_dag_runner_test.py b/tfx/orchestration/local/local_dag_runner_test.py
index c7199a0d1d..1e7a80379f 100644
--- a/tfx/orchestration/local/local_dag_runner_test.py
+++ b/tfx/orchestration/local/local_dag_runner_test.py
@@ -18,7 +18,6 @@
 from typing import Any, Dict, List
 
 import absl.testing.absltest
-import pytest
 from tfx import types
 from tfx.dsl.compiler import compiler
 from tfx.dsl.components.base import base_component
@@ -165,8 +164,6 @@ def _getTestPipelineIR(self) -> pipeline_pb2.Pipeline:  # pylint: disable=invali
     c = compiler.Compiler()
     return c.compile(test_pipeline)
 
-  @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 testRun(self):
     local_dag_runner.LocalDagRunner().run(self._getTestPipeline())
     self.assertEqual(_executed_components, [
@@ -174,8 +171,6 @@ def testRun(self):
         '_FakeComponent.d', '_FakeComponent.e'
     ])
 
-  @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 testPartialRun(self):
     local_dag_runner.LocalDagRunner().run(
         self._getTestPipeline(),
@@ -184,8 +179,6 @@ def testPartialRun(self):
         _executed_components,
         ['_FakeComponent.a', '_FakeComponent.b', '_FakeComponent.c'])
 
-  @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 testRunWithIR(self):
     local_dag_runner.LocalDagRunner().run_with_ir(self._getTestPipelineIR())
     self.assertEqual(_executed_components, [
@@ -193,8 +186,6 @@ def testRunWithIR(self):
         '_FakeComponent.d', '_FakeComponent.e'
     ])
 
-  @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 testPartialRunWithIR(self):
     pr_opts = pipeline_pb2.PartialRun()
     pr_opts.to_nodes.append('c')
diff --git a/tfx/orchestration/local/local_pipeline_test.py b/tfx/orchestration/local/local_pipeline_test.py
index 1ad12f7d6b..dd8203bf19 100644
--- a/tfx/orchestration/local/local_pipeline_test.py
+++ b/tfx/orchestration/local/local_pipeline_test.py
@@ -28,7 +28,6 @@
 from typing import Any, List
 
 import absl.testing.absltest
-import pytest
 
 from tfx import types
 from tfx.dsl.compiler import compiler
@@ -182,8 +181,6 @@ def _getTestPipelineIR(self) -> pipeline_pb2.Pipeline:
     c = compiler.Compiler()
     return c.compile(test_pipeline)
 
-  @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 testSimplePipelineRun(self):
     self.assertEqual(self.RAN_COMPONENTS, [])
 
@@ -191,8 +188,6 @@ def testSimplePipelineRun(self):
 
     self.assertEqual(self.RAN_COMPONENTS, ['Load', 'Train', 'Validate'])
 
-  @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 testSimplePipelinePartialRun(self):
     self.assertEqual(self.RAN_COMPONENTS, [])
 
@@ -202,8 +197,6 @@ def testSimplePipelinePartialRun(self):
 
     self.assertEqual(self.RAN_COMPONENTS, ['Load', 'Train'])
 
-  @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 testSimplePipelineRunWithIR(self):
     self.assertEqual(self.RAN_COMPONENTS, [])
 
@@ -211,8 +204,6 @@ def testSimplePipelineRunWithIR(self):
 
     self.assertEqual(self.RAN_COMPONENTS, ['Load', 'Train', 'Validate'])
 
-  @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 testSimplePipelinePartialRunWithIR(self):
     self.assertEqual(self.RAN_COMPONENTS, [])
 
diff --git a/tfx/orchestration/portable/inputs_utils_test.py b/tfx/orchestration/portable/inputs_utils_test.py
index c55cb20ec8..c077f518ce 100644
--- a/tfx/orchestration/portable/inputs_utils_test.py
+++ b/tfx/orchestration/portable/inputs_utils_test.py
@@ -14,7 +14,6 @@
 """Tests for tfx.orchestration.portable.inputs_utils."""
 import collections
 import os
-import pytest
 
 from tfx import types
 from tfx.dsl.compiler import placeholder_utils
@@ -147,8 +146,6 @@ def testResolveParametersFail(self):
     with self.assertRaisesRegex(RuntimeError, 'Parameter value not ready'):
       inputs_utils.resolve_parameters(parameters)
 
-  @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 testResolveInputArtifacts(self):
     pipeline = self.load_pipeline_proto(
         'pipeline_for_input_resolver_test.pbtxt')
@@ -254,8 +251,6 @@ def _setup_pipeline_for_input_resolver_test(self, num_examples=1):
     )
     self._examples = output_dict['output_examples']
 
-  @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 testResolveInputArtifacts_Normal(self):
     self._setup_pipeline_for_input_resolver_test()
 
@@ -266,8 +261,6 @@ def testResolveInputArtifacts_Normal(self):
     self.assertArtifactMapListEqual([{'examples_1': self._examples,
                                       'examples_2': self._examples}], 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)
   def testResolveInputArtifacts_FilterOutInsufficient(self):
     self._setup_pipeline_for_input_resolver_test()
     self._my_transform.inputs.inputs['examples_1'].min_count = 2
diff --git a/tfx/orchestration/portable/launcher_test.py b/tfx/orchestration/portable/launcher_test.py
index 916047b6a3..c75fd9043b 100644
--- a/tfx/orchestration/portable/launcher_test.py
+++ b/tfx/orchestration/portable/launcher_test.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 """Tests for tfx.orchestration.portable.launcher."""
 
-import pytest
 import contextlib
 import copy
 import os
@@ -490,8 +489,6 @@ def testLauncher_EmptyOptionalInputTriggersExecution(self):
           ],
       )
 
-  @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 testLauncher_PublishingNewArtifactsAndUseCache(self):
     # In this test case, there are two executions:
     # In the first one,trainer reads the fake upstream outputs and publish
@@ -578,8 +575,6 @@ def testLauncher_PublishingNewArtifactsAndUseCache(self):
           ],
       )
 
-  @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 testLauncher_CacheIsSupportedForNodeWithNoOutput(self):
     # Even though a node has no output at all, the launcher should treat the
     # second execution as CACHED as long as the cache context is the same.
@@ -639,8 +634,6 @@ def testLauncher_CacheIsSupportedForNodeWithNoOutput(self):
           ],
       )
 
-  @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 testLauncher_CacheDisabled(self):
     # In this test case, there are two executions:
     # In the first one,trainer reads the fake upstream outputs and publish
@@ -757,8 +750,6 @@ def testLauncher_CacheDisabled(self):
           ],
       )
 
-  @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 testLauncher_ReEntry(self):
     # Some executors or runtime environment may reschedule the launcher job
     # before the launcher job can publish any results of the execution to MLMD.
@@ -830,8 +821,6 @@ def create_test_launcher(executor_operators):
     execution_preparation_result = third_test_launcher._prepare_execution()
     self.assertFalse(execution_preparation_result.is_execution_needed)
 
-  @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 testLauncher_ToleratesDoubleCleanup(self):
     # Some executors or runtime environment may delete stateful_working_dir,
     # tmp_dir and unexpectedly. The launcher should handle such cases gracefully
@@ -895,8 +884,6 @@ def testLauncher_ToleratesDoubleCleanup(self):
           ],
       )
 
-  @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 testLauncher_ExecutionFailed(self):
     # In the case that the executor failed and raises an execption.
     # An Execution will be published.
@@ -916,8 +903,6 @@ def testLauncher_ExecutionFailed(self):
     with self.assertRaises(FakeError):
       _ = test_launcher.launch()
 
-  @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 testLauncher_ExecutionFailedViaReturnCode(self):
     # In the case that the executor failed and raises an execption.
     # An Execution will be published.
@@ -965,8 +950,6 @@ def testLauncher_ExecutionFailedViaReturnCode(self):
           ],
       )
 
-  @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 testLauncher_with_CustomDriver_NewSpan(self):
     self.reloadPipelineWithNewRunId()
     test_launcher = launcher.Launcher(
@@ -1019,8 +1002,6 @@ def testLauncher_with_CustomDriver_NewSpan(self):
           ],
       )
 
-  @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 testLauncher_with_CustomDriver_ExistingSpan(self):
     LauncherTest.fakeExampleGenOutput(self._mlmd_connection, self._example_gen,
                                       2, 1)
diff --git a/tfx/orchestration/portable/outputs_utils_test.py b/tfx/orchestration/portable/outputs_utils_test.py
index 0b643baeab..e14cb7d0b6 100644
--- a/tfx/orchestration/portable/outputs_utils_test.py
+++ b/tfx/orchestration/portable/outputs_utils_test.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 """Tests for tfx.orchestration.portable.output_utils."""
 
-import pytest
 import os
 from unittest import mock
 
@@ -251,8 +250,6 @@ def _get_external_uri_for_test(self, uri):
   @parameterized.parameters(
       (pipeline_pb2.Pipeline.SYNC, 'test_pipeline:test_run_0:test_node:1'),
       (pipeline_pb2.Pipeline.ASYNC, 'test_pipeline:test_node:1'))
-  @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 testGenerateOutputArtifacts(self, exec_mode, artifact_name_prefix):
     output_artifacts = self._output_resolver(
         exec_mode).generate_output_artifacts(1)
@@ -391,8 +388,6 @@ def testGetTmpDir(self):
     self.assertRegex(tmp_dir,
                      '.*/test_node/.system/executor_execution/1/.temp/')
 
-  @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 testMakeClearAndRemoveOutputDirs(self):
     output_artifacts = self._output_resolver().generate_output_artifacts(1)
     outputs_utils.make_output_dirs(output_artifacts)
@@ -415,8 +410,6 @@ def testMakeClearAndRemoveOutputDirs(self):
           continue
         self.assertFalse(fileio.exists(artifact.uri))
 
-  @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 testMakeOutputDirsArtifactAlreadyExists(self):
     output_artifacts = self._output_resolver().generate_output_artifacts(1)
     outputs_utils.make_output_dirs(output_artifacts)
@@ -442,8 +435,6 @@ def testMakeOutputDirsArtifactAlreadyExists(self):
           with fileio.open(os.path.join(artifact.uri, 'output'), 'r') as f:
             self.assertEqual(f.read(), 'test')
 
-  @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 testOmitLifeCycleManagementForExternalArtifact(self):
     """Test that it omits lifecycle management for external artifacts."""
     external_artifacts = self._output_resolver().generate_output_artifacts(1)
@@ -548,8 +539,6 @@ def testGetOrchestratorGeneratedBclDir(self):
     self.assertEqual(actual_bcl_dir, expected_bcl_dir)
     self.assertTrue(fileio.exists(actual_bcl_dir))
 
-  @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 testIntermediateArtifactState(self):
     pipeline_node = text_format.Parse(
         """
diff --git a/tfx/orchestration/portable/partial_run_utils_test.py b/tfx/orchestration/portable/partial_run_utils_test.py
index f54c50eb08..1fc9ddd005 100644
--- a/tfx/orchestration/portable/partial_run_utils_test.py
+++ b/tfx/orchestration/portable/partial_run_utils_test.py
@@ -14,7 +14,6 @@
 """Tests for tfx.orchestration.portable.partial_run_utils."""
 
 
-import pytest
 from collections.abc import Sequence
 from typing import Dict, List, Mapping, Optional, Set, Tuple, Union
 from unittest import mock
@@ -760,8 +759,6 @@ def assertResultEqual(self, pipeline_pb: pipeline_pb2.Pipeline,
         result_artifact.read()
         self.assertEqual(result_artifact.value, exp_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)
   def testArtifactRecyler_MultiplePipelines(self):
     """Tests that ArtifactRecyler works with multiple pipelines."""
     load = Load(start_num=1)
@@ -806,8 +803,6 @@ def testArtifactRecyler_MultiplePipelines(self):
           artifact_recyler._get_base_pipeline_run_context().name,
       )
 
-  @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 testSnapshot_removeFirstNode(self):
     """Tests that partial run with the first node removed works."""
     ############################################################################
@@ -912,8 +907,6 @@ def testSnapshot_removeFirstNode(self):
     ############################################################################
     self.assertResultEqual(pipeline_pb_run_2, 6)
 
-  @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 testReusePipelineArtifacts_twoIndependentSubgraphs(self):
     """Tests a sequence of partial runs with independent sub-graphs."""
     ############################################################################
@@ -1169,8 +1162,6 @@ def testReusePipelineArtifacts_twoIndependentSubgraphs(self):
               pipeline_run_contexts['run_3'], pipeline_run_contexts['run_4']
           ])
 
-  @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 testReusePipelineArtifacts_preventInconsistency(self):
     """Tests that a tricky sequence of partial runs raises an error."""
     ############################################################################
@@ -1366,8 +1357,6 @@ def testReusePipelineArtifacts_preventInconsistency(self):
     beam_dag_runner.BeamDagRunner().run_with_ir(pipeline_pb_run_5)
     self.assertResultEqual(pipeline_pb_run_5, 5)
 
-  @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 testNonExistentBaseRunId_lookupError(self):
     """Raise error if user provides non-existent base_run_id."""
     load = Load(start_num=1)
@@ -1391,8 +1380,6 @@ def testNonExistentBaseRunId_lookupError(self):
                                 'pipeline_run_id .* not found in MLMD.'):
       beam_dag_runner.BeamDagRunner().run_with_ir(pipeline_pb_run_2)
 
-  @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 testNonExistentNodeId_lookupError(self):
     """Raise error if user provides non-existent pipeline_run_id or node_id."""
     load = Load(start_num=1)
@@ -1417,8 +1404,6 @@ def testNonExistentNodeId_lookupError(self):
                                 'pipeline_run_id .* not found in MLMD.'):
       beam_dag_runner.BeamDagRunner().run_with_ir(pipeline_pb_run_2)
 
-  @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 testNoPreviousSuccessfulExecution_lookupError(self):
     """Raise error if user tries to reuse node w/o any successful Executions."""
     load_fail = LoadFail(start_num=1)
@@ -1443,8 +1428,6 @@ def testNoPreviousSuccessfulExecution_lookupError(self):
                                 'No previous successful executions found'):
       beam_dag_runner.BeamDagRunner().run_with_ir(pipeline_pb_run_2)
 
-  @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 testIdempotence_retryReusesRegisteredCacheExecution(self):
     """Ensures that there is only one registered cache execution.
 
@@ -1512,8 +1495,6 @@ def testIdempotence_retryReusesRegisteredCacheExecution(self):
               ]))
       self.assertLen(new_cache_executions, 1)
 
-  @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 testIdempotence_retryReusesPreviousSuccessfulCacheExecution(self):
     """Ensures idempotence.
 
@@ -1564,8 +1545,6 @@ def testIdempotence_retryReusesPreviousSuccessfulCacheExecution(self):
     beam_dag_runner.BeamDagRunner().run_with_ir(pipeline_pb_run_2)
     self.assertResultEqual(pipeline_pb_run_2, 6)
 
-  @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 testReusePipelineArtifacts_missingNewRunId_error(self):
     """If pipeline IR has no run id, and user does not provide it, fail."""
     ############################################################################
@@ -1636,8 +1615,6 @@ def testReusePipelineArtifacts_missingNewRunId_error(self):
     beam_dag_runner.BeamDagRunner().run_with_ir(pipeline_pb_run_2)
     self.assertResultEqual(pipeline_pb_run_2, 6)
 
-  @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 testReusePipelineArtifacts_inconsistentNewRunId_error(self):
     """If pipeline IR's run_id differs from user-provided run_id, fail."""
     ############################################################################
@@ -1698,8 +1675,6 @@ def testReusePipelineArtifacts_inconsistentNewRunId_error(self):
             m, pipeline_pb_run_2, base_run_id='run_1',
             new_run_id='run_3')  # <-- user error here
 
-  @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 testReusePipelineArtifacts_SeparateBranches(self):
     """Tests partial run with separate branches."""
     ############################################################################
diff --git a/tfx/orchestration/python_execution_binary/python_execution_binary_utils_test.py b/tfx/orchestration/python_execution_binary/python_execution_binary_utils_test.py
index 285074b898..45b09a90eb 100644
--- a/tfx/orchestration/python_execution_binary/python_execution_binary_utils_test.py
+++ b/tfx/orchestration/python_execution_binary/python_execution_binary_utils_test.py
@@ -15,7 +15,6 @@
 
 
 
-import pytest
 from typing import Dict, List, Union
 
 import tensorflow as tf
@@ -45,8 +44,6 @@ class _MyArtifact(artifact.Artifact):
   }
 
 
-@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 PythonExecutorBinaryUtilsTest(tf.test.TestCase):
 
   def _convert_to_artifact_proto(
diff --git a/tfx/types/artifact_test.py b/tfx/types/artifact_test.py
index b8af072bc5..b7e6eb2b38 100644
--- a/tfx/types/artifact_test.py
+++ b/tfx/types/artifact_test.py
@@ -15,6 +15,8 @@
 
 import gc
 import json
+import importlib
+import pytest
 import textwrap
 from unittest import mock
 
@@ -30,6 +32,12 @@
 from ml_metadata.proto import metadata_store_pb2
 
 
+@pytest.fixture(scope="module", autouse=True)
+def cleanup():
+  yield
+  importlib.reload(struct_pb2)
+
+
 Dataset = system_artifacts.Dataset
 
 
diff --git a/tfx/utils/doc_controls_test.py b/tfx/utils/doc_controls_test.py
index 6a2d5f2c2a..c096174fae 100644
--- a/tfx/utils/doc_controls_test.py
+++ b/tfx/utils/doc_controls_test.py
@@ -15,15 +15,12 @@
 
 
 
-import pytest
 import tensorflow as tf
 
 from tfx.utils import doc_controls as tfx_doc_controls
 from tensorflow.tools.docs import doc_controls  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
 
 
-@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 DocControlsTest(tf.test.TestCase):
 
   def testDocControls(self):
@@ -33,6 +30,8 @@ def testDocControls(self):
                      doc_controls.do_not_doc_in_subclasses)
 
   def testDocumentSuccess(self):
+    # Clean up EXTRA_DOCS since pytest can import other modules in other tests.
+    tfx_doc_controls.EXTRA_DOCS = dict()
     documented_test_key = tfx_doc_controls.documented('test key', 'test value')
     self.assertEqual(1, len(tfx_doc_controls.EXTRA_DOCS))
     self.assertEqual('test value',