code: def test_bazel_sharding_two_shards(mock_build_dir):
      """Test e2e working of the script with sharding.
  
      Assert that the two shards are balanced as expected.
      """
      output_1_list = bazel_sharding.main(["..."], index=0, count=2)
      output_1 = set(output_1_list)
  
      output_2_list = bazel_sharding.main(["..."], index=1, count=2)
      output_2 = set(output_2_list)
  
      assert output_1.union(output_2) == set(all_rules)
  
      # We should be deterministic, therefore we can hardcode this
      assert output_1_list == [
          f"//{WORKSPACE_KEY}:test_eternal",
          f"//{WORKSPACE_KEY}:test_large",
          f"//{WORKSPACE_KEY}:test_long",
          f"//{WORKSPACE_KEY}:test_small",
      ]
      assert output_2_list == [
          f"//{WORKSPACE_KEY}:test_both_size_and_timeout",
          f"//{WORKSPACE_KEY}:test_default",
          f"//{WORKSPACE_KEY}:test_enormous",
          f"//{WORKSPACE_KEY}:test_medium",
          f"//{WORKSPACE_KEY}:test_moderate",
          f"//{WORKSPACE_KEY}:test_short",
      ]
  
      output_1_naive_list = bazel_sharding.main(
          ["..."], index=0, count=2, sharding_strategy="naive"
      )
      output_1_naive = set(output_1_naive_list)
  
      output_2_naive_list = bazel_sharding.main(
          ["..."], index=1, count=2, sharding_strategy="naive"
      )
      output_2_naive = set(output_2_naive_list)
  
      assert output_1_naive.union(output_2_naive) == set(all_rules)
      # We should be deterministic, therefore we can hardcode this
      assert output_1_naive_list == [
          f"//{WORKSPACE_KEY}:test_both_size_and_timeout",
          f"//{WORKSPACE_KEY}:test_enormous",
          f"//{WORKSPACE_KEY}:test_large",
          f"//{WORKSPACE_KEY}:test_medium",
          f"//{WORKSPACE_KEY}:test_short",
      ]
      assert output_2_naive_list == [
          f"//{WORKSPACE_KEY}:test_default",
          f"//{WORKSPACE_KEY}:test_eternal",
          f"//{WORKSPACE_KEY}:test_long",
          f"//{WORKSPACE_KEY}:test_moderate",
          f"//{WORKSPACE_KEY}:test_small",
      ]
  
file_path: ray/ci/ray_ci/test_bazel_sharding.py
imports: from typing import List
  import pytest
  import os
  import shutil
  import sys
  import tempfile
  import bazel_sharding
label_desc: Use image node_id 7 for calling a global function with example usage: test_bazel_sharding_two_shards(mock_build_dir) without return types