Unittest (PyTest) function used in mapPartitions() in PySpark

I have a use case where in I'm writing a dataframe to DynamoDB.

The issue is, that although I'm able to run the code on the cluster, the function used inside of mapPartitions() is getting hard to test.

Here is my code

import boto3
import json

def write_df(table_name, region):
    the_list=[]

    def my_map_partition(partition):
        table = boto3.resource("dynamodb", region=region).Table(table_name)
        for row in partition:
            table.put_item(Item = json.loads(row.col1))
            the_list.append(row.col1)
        return the_list

    return my_map_partition

def load_to_db(df):
    df.rdd.mapPartitions(write_df("actual_table", "eu-west-2"))

@mock_dynamodb2
def init_dynamo():
    client = boto3.client("dynamodb", "eu-west-2")
    client.create_table(
      TableName="my_test_table",
      KeySchema=[
        {"AttributeName": "partitionKey", "KeyType": "HASH"},
        {"AttributeName": "sortKey", "KeyType": "RANGE"}
      ],
      AttributeDefinitions=[
        {"AttributeName": "partitionKey", "AttributeType": "S"},
        {"AttributeName": "sortKey", "AttributeType": "S"},
      ]
    )

def test_write_df():
    self.init_dynamo()

    columns = ['col1', 'col2']
    vals = [('{"partitionKey":"pk1", "sortKey":"sk1", "random_value":"hello_world"}', 'dummy'),('{"partitionKey":"pk2", "sortKey":"sk2", "random_value":"hello_world"}', 'dummy')]
    data = spark.createDataframe(vals, columns).select("col1")

    a = data.rdd.mapPartitions(write_df("my_test_table", "eu-west-2"))
    a.count()

    pass #here is where I will assert values

Here is the error

objc[99130]: +[__NSCFConstantString initialize] may have been in progress in another thread when fork() was called.
objc[99130]: +[__NSCFConstantString initialize] may have been in progress in another thread when fork() was called. We cannot safely call it or ignore it in the fork() child process. Crashing instead. Set a breakpoint on objc_initializeAfterForkError to debug.
22/07/15 14:24:38 ERROR Executor: Exception in task 0.0 in stage 3.0 (TID 32)
org.apache.spark.SparkException: Python worker exited unexpectedly (crashed)
    at org.apache.spark.api.python.BasePythonRunner$ReaderIterator$$anonfun$1.applyOrElse(PythonRunner.scala:585)
    at org.apache.spark.api.python.BasePythonRunner$ReaderIterator$$anonfun$1.applyOrElse(PythonRunner.scala:567)
.
.
    ERROR TaskSetManager: Task 0 in stage 3.0 failed 1 times; aborting job
    
    _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
    /usr/lib/python3.9/site-packages/pyspark/sql/session.py:66: in toDF
        return sparkSession.createDataFrame(self, schema, sampleRatio)
    /usr/lib/python3.9/site-packages/pyspark/sql/session.py:675: in createDataFrame
        return self._create_dataframe(data, schema, samplingRatio, verifySchema)
    /usr/lib/python3.9/site-packages/pyspark/sql/session.py:698: in _create_dataframe
        rdd, schema = self._createFromRDD(data.map(prepare), schema, samplingRatio)
    /usr/lib/python3.9/site-packages/pyspark/sql/session.py:486: in _createFromRDD
        struct = self._inferSchema(rdd, samplingRatio, names=schema)
    /usr/lib/python3.9/site-packages/pyspark/sql/session.py:460: in _inferSchema
        first = rdd.first()
    /usr/lib/python3.9/site-packages/pyspark/rdd.py:1588: in first
        rs = self.take(1)
    /usr/lib/python3.9/site-packages/pyspark/rdd.py:1568: in take
        res = self.context.runJob(self, takeUpToNumLeft, p)
    /usr/lib/python3.9/site-packages/pyspark/context.py:1227: in runJob
        sock_info = self._jvm.PythonRDD.runJob(self._jsc.sc(), mappedRDD._jrdd, partitions)
    /usr/lib/python3.9/site-packages/py4j/java_gateway.py:1309: in __call__
        return_value = get_return_value(
    /usr/lib/python3.9/site-packages/pyspark/sql/utils.py:111: in deco
        return f(*a, **kw)
    _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
    
    answer = 'xro77'
    gateway_client = <py4j.clientserver.JavaClient object at 0x124798820>
    target_id = 'z:org.apache.spark.api.python.PythonRDD', name = 'runJob'
    
        def get_return_value(answer, gateway_client, target_id=None, name=None):
            """Converts an answer received from the Java gateway into a Python object.
            For example, string representation of integers are converted to Python
            integer, string representation of objects are converted to JavaObject
            instances, etc.
        
            :param answer: the string returned by the Java gateway
            :param gateway_client: the gateway client used to communicate with the Java
                Gateway. Only necessary if the answer is a reference (e.g., object,
                list, map)
            :param target_id: the name of the object from which the answer comes from
                (e.g., *object1* in `object1.hello()`). Optional.
            :param name: the name of the member from which the answer comes from
                (e.g., *hello* in `object1.hello()`). Optional.
            """
            if is_error(answer)[0]:
                if len(answer) > 1:
                    type = answer[1]
                    value = OUTPUT_CONVERTER[type](answer[2:], gateway_client)
                    if answer[1] == REFERENCE_TYPE:
    >                   raise Py4JJavaError(
                            "An error occurred while calling {0}{1}{2}.\n".
                            format(target_id, ".", name), value)
    E                   py4j.protocol.Py4JJavaError: An error occurred while calling 

I'm not sure what is wrong in the unit tests since this is a working code.

Can someone guide me in the right direction on how do I test a function in PySpark that is being used for mapPartitions() ?