[Data][CI] Add lint to ensure tests actually run

.buildkite/lint.rayci.yml

@@ -13,7 +13,7 @@ steps:
       - clang_format
       - code_format
       - pre_commit
-      - untested_code_snippet
+      - semgrep_lint
       - banned_words
       - doc_readme
       - dashboard_format

ci/lint/lint.sh

@@ -60,7 +60,7 @@ code_format() {
   FORMAT_SH_PRINT_DIFF=1 ./ci/lint/format.sh --all-scripts
 }
 
-untested_code_snippet() {
+semgrep_lint() {
   pip install -c python/requirements_compiled.txt semgrep
   semgrep ci --config semgrep.yml
 }

python/ray/data/tests/test_arrow_serialization.py

@@ -653,3 +653,11 @@ def test_variable_shape_tensor_serialization():
     payload = PicklableArrayPayload.from_array(ar)
     ar2 = payload.to_array()
     assert ar == ar2
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/test_block.py

@@ -156,3 +156,11 @@ def test_find_partitions_duplicates():
     assert partitions[1].to_pydict() == {"value": []}  # [1,2)
     assert partitions[2].to_pydict() == {"value": [2, 2, 2, 2, 2]}  # [2,3)
     assert partitions[3].to_pydict() == {"value": []}  # >=3
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/test_block_boundaries.py

@@ -60,3 +60,11 @@ def test_groupby_map_groups_get_block_boundaries_with_nan():
     )
 
     assert list(indices) == [0, 1, 2, 4, 6, 7]
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/test_expression_evaluator.py

@@ -338,3 +338,11 @@ def test_filter_bad_expression(sample_data):
     sample_data_path, _ = sample_data
     with pytest.raises(pa.ArrowInvalid):
         pq.read_table(sample_data_path, filters=filters)
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/test_hash_shuffle.py

@@ -491,3 +491,11 @@ def test_hash_shuffle_operator_remote_args(
                 op._aggregator_pool._aggregator_ray_remote_args
                 == tc.expected_ray_remote_args
             )
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/test_operator_fusion.py

@@ -729,3 +729,11 @@ def test_zero_copy_fusion_eliminate_build_output_blocks(
             BatchMapTransformFn,
         ],
     )
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/data/tests/test_ref_bundle.py

@@ -52,3 +52,11 @@ def _get_obj_locs(obj_refs):
         "2": 7168,  # first_block_ref, second_block_ref, third_block_ref
         "3": 3072,  # first_block_ref, second_block_ref
     } == preferred_object_locs
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

python/requirements/lint-requirements.txt

@@ -5,7 +5,7 @@ mypy==1.7.0
 mypy-extensions==1.0.0
 types-PyYAML==6.0.12.2
 black==22.10.0
-semgrep==1.32.0
+semgrep==1.136.0
 shellcheck-py==0.7.1.1
 yq
 types-pycurl==7.45.3.20240421

semgrep.yml

@@ -27,3 +27,26 @@ rules:
     message: "Don't use 'code-block:: python', it's not tested! Use 'testcode' instead! For more information, see https://docs.ray.io/en/master/ray-contribute/writing-code-snippets.html."
     pattern: "code-block:: python"
     severity: ERROR
+
+  - id: missing-pytest-main
+    paths:
+      include:
+        - "python/ray/data/tests/**/test_*.py"
+    languages:
+      - python
+    message: |
+      Add the following snippet to the end of the file so that the tests run in CI.
+
+      ```python
+      if __name__ == "__main__":
+          import sys
+
+          import pytest
+
+          sys.exit(pytest.main(["-v", __file__]))
+      ```
+    patterns:
+      - pattern-regex: |
+          (?s)(.*)
+      - pattern-not-regex: pytest.main
+    severity: ERROR