From b498e4ab3bbbdd57c9fe56f7c8b6c844970f9311 Mon Sep 17 00:00:00 2001
From: Anand Inguva <anandinguva98@gmail.com>
Date: Wed, 13 Dec 2023 01:53:47 -0500
Subject: [PATCH] Skip few tests

---
 .../testing/benchmarks/cloudml/cloudml_benchmark_test.py   | 5 ++++-
 .../benchmarks/{cloudml => }/mltransform/__init__.py       | 0
 .../testing/benchmarks/{cloudml => }/mltransform/criteo.py | 0
 .../benchmarks/{cloudml => }/mltransform/criteo_test.py    | 7 +------
 4 files changed, 5 insertions(+), 7 deletions(-)
 rename sdks/python/apache_beam/testing/benchmarks/{cloudml => }/mltransform/__init__.py (100%)
 rename sdks/python/apache_beam/testing/benchmarks/{cloudml => }/mltransform/criteo.py (100%)
 rename sdks/python/apache_beam/testing/benchmarks/{cloudml => }/mltransform/criteo_test.py (89%)

diff --git a/sdks/python/apache_beam/testing/benchmarks/cloudml/cloudml_benchmark_test.py b/sdks/python/apache_beam/testing/benchmarks/cloudml/cloudml_benchmark_test.py
index ccaf2681c2d3..905aa18fb5d7 100644
--- a/sdks/python/apache_beam/testing/benchmarks/cloudml/cloudml_benchmark_test.py
+++ b/sdks/python/apache_beam/testing/benchmarks/cloudml/cloudml_benchmark_test.py
@@ -55,9 +55,9 @@ def _publish_metrics(pipeline, metric_value, metrics_table, metric_name):
   )])
 
 
-@unittest.skip('Remve this line to run this test')
 @pytest.mark.uses_tft
 class CloudMLTFTBenchmarkTest(unittest.TestCase):
+  @unittest.skip('Remve this line to run this test')
   def test_cloudml_benchmark_criteo_small(self):
     test_pipeline = TestPipeline(is_integration_test=True)
     extra_opts = {}
@@ -79,6 +79,7 @@ def test_cloudml_benchmark_criteo_small(self):
         metrics_table=metrics_table,
         metric_name='runtime_sec')
 
+  @unittest.skip('Remve this line to run this test')
   def test_cloudml_benchmark_cirteo_no_shuffle_10GB(self):
     test_pipeline = TestPipeline(is_integration_test=True)
     extra_opts = {}
@@ -112,6 +113,8 @@ def test_cloudml_benchmark_criteo_10GB(self):
     extra_opts['output'] = os.path.join(
         _OUTPUT_GCS_BUCKET_ROOT, uuid.uuid4().hex)
     start_time = time.time()
+    extra_opts['job_name'] = (
+        'tft-criteo-dataset-{}-10'.format(uuid.uuid4().hex))
     workflow.run(test_pipeline.get_full_options_as_args(**extra_opts))
     end_time = time.time()
 
diff --git a/sdks/python/apache_beam/testing/benchmarks/cloudml/mltransform/__init__.py b/sdks/python/apache_beam/testing/benchmarks/mltransform/__init__.py
similarity index 100%
rename from sdks/python/apache_beam/testing/benchmarks/cloudml/mltransform/__init__.py
rename to sdks/python/apache_beam/testing/benchmarks/mltransform/__init__.py
diff --git a/sdks/python/apache_beam/testing/benchmarks/cloudml/mltransform/criteo.py b/sdks/python/apache_beam/testing/benchmarks/mltransform/criteo.py
similarity index 100%
rename from sdks/python/apache_beam/testing/benchmarks/cloudml/mltransform/criteo.py
rename to sdks/python/apache_beam/testing/benchmarks/mltransform/criteo.py
diff --git a/sdks/python/apache_beam/testing/benchmarks/cloudml/mltransform/criteo_test.py b/sdks/python/apache_beam/testing/benchmarks/mltransform/criteo_test.py
similarity index 89%
rename from sdks/python/apache_beam/testing/benchmarks/cloudml/mltransform/criteo_test.py
rename to sdks/python/apache_beam/testing/benchmarks/mltransform/criteo_test.py
index a3e45962462d..cf2a1ef27acc 100644
--- a/sdks/python/apache_beam/testing/benchmarks/cloudml/mltransform/criteo_test.py
+++ b/sdks/python/apache_beam/testing/benchmarks/mltransform/criteo_test.py
@@ -27,17 +27,15 @@
 
 try:
   import apache_beam.testing.benchmarks.cloudml.cloudml_benchmark_constants_lib as constants
-  #   from apache_beam.examples.ml_transform import vocab_tfidf_processing
   from apache_beam.testing.load_tests.load_test_metrics_utils import InfluxDBMetricsPublisherOptions
   from apache_beam.testing.load_tests.load_test_metrics_utils import MetricsReader
   from apache_beam.testing.test_pipeline import TestPipeline
-  from apache_beam.examples.ml_transform import criteo
+  from apache_beam.testing.benchmarks.mltransform import criteo
 except ImportError:  # pylint: disable=bare-except
   raise unittest.SkipTest('tensorflow_transform is not installed.')
 
 _INPUT_GCS_BUCKET_ROOT = 'gs://apache-beam-ml/datasets/cloudml/criteo'
 _OUTPUT_GCS_BUCKET_ROOT = 'gs://temp-storage-for-end-to-end-tests/tft/'
-_DISK_SIZE = 150
 
 
 @pytest.mark.uses_tft
@@ -54,9 +52,6 @@ def test_process_criteo_10GB_dataset(self):
 
     extra_opts['frequency_threshold'] = 0
 
-    # dataflow pipeliens options
-    extra_opts['disk_size_gb'] = _DISK_SIZE
-    extra_opts['machine_type'] = 'e2-highmem-2'
     extra_opts['job_name'] = (
         'mltransform-criteo-dataset-{}-10'.format(uuid.uuid4().hex))
     # start_time = time.time()