Re: Unit Testing with PyTest in Databricks - Modul...

saurabh18cs · ‎10-25-2024

Hi,

After trying a lot I could able to see some success , see if this is what you all are looking for :

notebook_test.py (this is python code file)

from pyspark.sql import functions as F

def sum_values(df😞

return df.agg(F.sum("value")).first()[0]

def reverse(s😞

return s[::-1]

# Return the functions as a dictionary

# dbutils.notebook.exit({

# "sum_values": sum_values,

# "reverse": reverse

# })

test_sum (this is notebook , both sitting parallel to each other)

cmd1

!pip install pytest

cmd2

import pytest

import os

import sys

sys.dont_write_bytecode = True

os.chdir("/Workspace/Users/saurabh.............../")

cmd3

import pytest

from pyspark.sql import SparkSession

from notebook_test import sum_values, reverse

# Run the notebook and import the functions

# notebook_path = "/Workspace/Users/saurabh................../notebook_test1"

# notebook_output = dbutils.notebook.run(notebook_path, 60)

# functions = eval(notebook_output)

# sum_values = functions["sum_values"]

# reverse = functions["reverse"]

@pytest.fixture(scope="module")

def spark():

spark = SparkSession.builder \

.appName("pytest-pyspark-local-testing") \

.master("local[*]") \

.getOrCreate()

yield spark

spark.stop()

def test_sum_values(spark😞

data = [(1,), (2,), (3,)]

df = spark.createDataFrame(data, ["value"])

result = sum_values(df)

assert result == 6

def test_reverse():

assert reverse("hello") == "olleh"

assert reverse("world") == "dlrow"

assert reverse("") == ""

assert reverse("a") == "a"

cmd4

# In Databricks notebook

pytest.main(["-v"], plugins = [test_sum_values(spark)])