How to unittest pyspark `withColumn` action - Python 3?

user_10581944 0
Viewed 138

I'm learning pyspark, I have a function:

import re

def function_1(string):
    new_string = re.sub(r"!", " ", string)
    return new_string


udf_function_1 = udf(lambda s: function_1(s), StringType())


def function_2(data):
    new_data = data \
        .withColumn("column_1", udf_function_1("column_1"))
    return new_data

My question is how to write unittest for function_2() in Python.

pyspark apache-spark-sql python-unittest unit-testing python
1 Answers

what do you exactly want to test in function_2?

Below is a simple test saved in a file called sample_test.py. I used pytest but you can right very similar code in unittest.

# sample_test.py

from pyspark import sql


spark = sql.SparkSession.builder \
        .appName("local-spark-session") \
        .getOrCreate()
        
def test_create_session():
    assert isinstance(spark, sql.SparkSession) == True
    assert spark.sparkContext.appName == 'local-spark-session'

def test_spark_version():
    assert spark.version == '3.1.2'

running the test...

C:\Users\user\Desktop>pytest -v sample_test.py
============================================= test session starts =============================================
platform win32 -- Python 3.6.7, pytest-6.2.5, py-1.10.0, pluggy-1.0.0 -- c:\users\user\appdata\local\programs\python\python36\python.exe
cachedir: .pytest_cache
rootdir: C:\Users\user\Desktop
collected 2 items

sample_test.py::test_create_session PASSED                                                               [ 50%]
sample_test.py::test_spark_version PASSED                                                                [100%]

============================================== 2 passed in 4.81s ==============================================
user_83133360
Related