diff --git a/.buildkite/lint.rayci.yml b/.buildkite/lint.rayci.yml
index 9762ce7a8261..368b81b036b4 100644
--- a/.buildkite/lint.rayci.yml
+++ b/.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
diff --git a/.semgrepignore b/.semgrepignore
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/ci/lint/lint.sh b/ci/lint/lint.sh
index 53d7156e86c8..304429cb85d6 100755
--- a/ci/lint/lint.sh
+++ b/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
 }
diff --git a/python/ray/data/tests/test_arrow_serialization.py b/python/ray/data/tests/test_arrow_serialization.py
index 0a437920f908..014c28dc49ed 100644
--- a/python/ray/data/tests/test_arrow_serialization.py
+++ b/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__]))
diff --git a/python/ray/data/tests/test_block.py b/python/ray/data/tests/test_block.py
index f897a8bc4529..b7c587902f2c 100644
--- a/python/ray/data/tests/test_block.py
+++ b/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__]))
diff --git a/python/ray/data/tests/test_block_boundaries.py b/python/ray/data/tests/test_block_boundaries.py
index 7f5d5be7a72d..6fd40ba95185 100644
--- a/python/ray/data/tests/test_block_boundaries.py
+++ b/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__]))
diff --git a/python/ray/data/tests/test_expression_evaluator.py b/python/ray/data/tests/test_expression_evaluator.py
index 0754539e0829..1e8b50a83928 100644
--- a/python/ray/data/tests/test_expression_evaluator.py
+++ b/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__]))
diff --git a/python/ray/data/tests/test_hash_shuffle.py b/python/ray/data/tests/test_hash_shuffle.py
index 6d9909799756..e443610e589c 100644
--- a/python/ray/data/tests/test_hash_shuffle.py
+++ b/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__]))
diff --git a/python/ray/data/tests/test_operator_fusion.py b/python/ray/data/tests/test_operator_fusion.py
index 28db544c9f04..abc9b5607ab2 100644
--- a/python/ray/data/tests/test_operator_fusion.py
+++ b/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__]))
diff --git a/python/ray/data/tests/test_ref_bundle.py b/python/ray/data/tests/test_ref_bundle.py
index 549bbad311f5..3bd83349f7fb 100644
--- a/python/ray/data/tests/test_ref_bundle.py
+++ b/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__]))
diff --git a/python/ray/data/tests/test_should_fail.py b/python/ray/data/tests/test_should_fail.py
new file mode 100644
index 000000000000..108caad3e6bd
--- /dev/null
+++ b/python/ray/data/tests/test_should_fail.py
@@ -0,0 +1,2 @@
+def test_ham():
+    assert False
diff --git a/python/requirements/lint-requirements.txt b/python/requirements/lint-requirements.txt
index 32cd3564f168..68bdd78433d4 100644
--- a/python/requirements/lint-requirements.txt
+++ b/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
diff --git a/semgrep.yml b/semgrep.yml
index 43b19533a019..be17af1bdeef 100644
--- a/semgrep.yml
+++ b/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