Fixed formatting issues and added unit test

Signed-off-by: Atharva <atharvakulkarni172003@gmail.com>

Author: Atharva1723

flytekit/bin/entrypoint.py

@@ -470,9 +470,9 @@ def setup_execution(
     if checkpoint_path is not None:
         checkpointer = SyncCheckpoint(checkpoint_dest=checkpoint_path, checkpoint_src=prev_checkpoint)
         logger.debug(f"Checkpointer created with source {prev_checkpoint} and dest {checkpoint_path}")
-    
+
     node_index = _compute_array_job_index()
-    
+
     execution_parameters = ExecutionParameters(
         execution_id=_identifier.WorkflowExecutionIdentifier(
             project=exe_project,

tests/flytekit/unit/bin/test_python_entrypoint.py

@@ -37,6 +37,7 @@
 from flytekit.models import literals as _literal_models
 from flytekit.models.core import errors as error_models, execution
 from flytekit.models.core import execution as execution_models
+from flytekit.bin.entrypoint import _compute_array_job_index
 from flytekit.core.utils import write_proto_to_file
 from flytekit.models.types import LiteralType, SimpleType
 
@@ -406,6 +407,16 @@ def test_dispatch_execute_system_error(mock_write_to_file, mock_upload_dir, mock
         assert "some system exception" in ed.error.message
         assert ed.error.origin == execution_models.ExecutionError.ErrorKind.SYSTEM
 
+@pytest.fixture
+def flyte_context():
+    """Fixture to set up a mock Flyte context."""
+    with mock.patch.object(context_manager.FlyteContext, 'current_context', return_value=mock.Mock()):
+        yield
+
+def test_compute_array_job_index(flyte_context):
+    assert _compute_array_job_index() == 0  
+    assert _compute_array_job_index(index=1) == 1 
+    assert _compute_array_job_index(index=2) == 2  
 
 def test_setup_disk_prefix():
     with setup_execution("qwerty") as ctx: