[SPARK-39179][PYTHON][TESTS] Improve the test coverage for pyspark/shuffle.py
[SPARK-39179][PYTHON][TESTS] Improve the test coverage for pyspark/shuffle.py

### What changes were proposed in this pull request?
This PR add test cases for shuffle.py

### Why are the changes needed?
To cover corner test cases and increase coverage. This will increase the coverage of shuffle.py to close to 90%

### Does this PR introduce _any_ user-facing change?
No - test only

### How was this patch tested?
CI in this PR should test it out

Closes #36701 from pralabhkumar/rk_test_taskcontext.

Authored-by: pralabhkumar <a99b3d4dcf6225dc035bc084a130a93c7225bad5@gmail.com>
Signed-off-by: Dongjoon Hyun <6ff48b33cbb828151035af58a92f18a39f1d1036@apache.org>

rednaxelafx/apache-spark,holdenk/spark,JoshRosen/spark,witgo/spark,cloud-fan/spark,holdenk/spark,WeichenXu123/spark,shaneknapp/spark,shaneknapp/spark,rednaxelafx/apache-spark,gengliangwang/spark,holdenk/spark,dongjoon-hyun/spark,witgo/spark,dongjoon-hyun/spark,vinodkc/spark,dongjoon-hyun/spark,hvanhovell/spark,vinodkc/spark,shaneknapp/spark,vinodkc/spark,zero323/spark,zzcclp/spark,hvanhovell/spark,WeichenXu123/spark,hvanhovell/spark,WeichenXu123/spark,srowen/spark,nchammas/spark,mahak/spark,mahak/spark,JoshRosen/spark,vinodkc/spark,JoshRosen/spark,JoshRosen/spark,dongjoon-hyun/spark,vinodkc/spark,cloud-fan/spark,HyukjinKwon/spark,zero323/spark,witgo/spark,apache/spark,zzcclp/spark,rednaxelafx/apache-spark,rednaxelafx/apache-spark,nchammas/spark,HyukjinKwon/spark,vinodkc/spark,JoshRosen/spark,shaneknapp/spark,mahak/spark,nchammas/spark,apache/spark,mahak/spark,apache/spark,srowen/spark,shaneknapp/spark,dongjoon-hyun/spark,HyukjinKwon/spark,rednaxelafx/apache-spark,cloud-fan/spark,wangyum/spark,cloud-fan/spark,WeichenXu123/spark,zzcclp/spark,zero323/spark,hvanhovell/spark,gengliangwang/spark,witgo/spark,zzcclp/spark,wangyum/spark,wangyum/spark,shaneknapp/spark,zero323/spark,mahak/spark,WeichenXu123/spark,WeichenXu123/spark,mahak/spark,ueshin/apache-spark,ueshin/apache-spark,dongjoon-hyun/spark,ueshin/apache-spark,ueshin/apache-spark,dongjoon-hyun/spark,shaneknapp/spark,apache/spark,srowen/spark,apache/spark,zero323/spark,witgo/spark,HyukjinKwon/spark,gengliangwang/spark,ueshin/apache-spark,srowen/spark,JoshRosen/spark,nchammas/spark,zero323/spark,WeichenXu123/spark,ueshin/apache-spark,holdenk/spark,zzcclp/spark,gengliangwang/spark,HyukjinKwon/spark,wangyum/spark,rednaxelafx/apache-spark,HyukjinKwon/spark,witgo/spark,wangyum/spark,apache/spark,gengliangwang/spark,holdenk/spark,ueshin/apache-spark,nchammas/spark,wangyum/spark,hvanhovell/spark,cloud-fan/spark,gengliangwang/spark,holdenk/spark,wangyum/spark,HyukjinKwon/spark,cloud-fan/spark,hvanhovell/spark,witgo/spark,hvanhovell/spark,nchammas/spark,mahak/spark,gengliangwang/spark,cloud-fan/spark,holdenk/spark,JoshRosen/spark,nchammas/spark,srowen/spark,zzcclp/spark,srowen/spark,srowen/spark,rednaxelafx/apache-spark,zero323/spark,apache/spark,vinodkc/spark,zzcclp/spark