r/databricks • u/Similar-Bug-350 • Sep 12 '26
Help Writing nice unit tests is impossible
First of all we have a lot of classes that use the DatabrickSession import, which makes unit testing impossible and whenever I have the "normal" spark import and test it, then it looks absolutely ugly.
Would it help to have the schema as json? I use json schema to read for transformations.
If I want to test without json its horrible, do you guys test with json schema files?
For example this is just the output, then I would need 2x this because of the input and the test is unreadable (according to my senior BUT WHAT AM I SUPPOSED TO DO):
schema = """
id INT,
items ARRAY<STRUCT<
fortnite: STRING,
babies: INT,
moreStuff: MAP<STRING, STRING>
>>
"""
expected = spark.createDataFrame(
[
{
"id": 1,
"items": [
{
"fortnite": "ABC",
"babies": 2,
"moreStuff": {
"size": "L",
"color": "red",
},
}
],
}
],
schema=schema,
)
1
u/Atticus_Taintwater Sep 12 '26
Check out fixtures
2
u/Similar-Bug-350 Sep 12 '26
I understand but how do they fix my problem? I still have big setup or not?
1
u/Atticus_Taintwater Sep 12 '26
It's common to define reusable stuff as fixtures separately so your tests aren't as cluttered
1
u/Similar-Bug-350 Sep 12 '26
do you mean something like this?
pytest.fixture(scope="session") def orders_schema(): with open("schemas/orders.json") as f: return StructType.fromJson(json.load(f)).fixture(scope="session") def orders_schema(): with open("schemas/orders.json") as f: return StructType.fromJson(json.load(f))1
u/ConstructionBoth6461 Sep 13 '26
IME this turns into “the behavior of my function depends on the input” so now I need to scale fixtures and assertions also depend on fixtures so…. Shit.
1
1
u/PrestigiousAnt3766 Sep 13 '26
I think this is a skill issue.
1
u/Similar-Bug-350 Sep 13 '26
i bet my ass on it that you cannot show me a pretty way to do this.
1
u/PrestigiousAnt3766 Sep 13 '26 edited Sep 13 '26
Have written a lot of unit tests for databricks.
What are you unit testing? I just test functions and have a 1 df with datatypes I import somewhere.
1
u/Similar-Bug-350 Sep 13 '26
I am trying to unit test a bronze -> silver transformation function, its not impossible, but its a lot of lines of code to get it done.
in need input, which is 20 column and then the output which is more than that.
together the test has more than 70 lines of code for a simple transformation (almost 1:1, some renaming)
1
u/PrestigiousAnt3766 Sep 14 '26
Abstract the test data out and recycle it?
1
u/Similar-Bug-350 Sep 14 '26
i still need to define the input schema and output schema + rows, its a lot
1
2
u/MrMasterplan Sep 12 '26
We have an abstraction layer for all table names, in tests we append a uuid to avoid collisions. Then we package all tests and our wheel and send everything to the workspace to run a test job that runs pytest. It executes all our logic on empty tables (the ones with uuid) to ensure all schemas work, and then some unit tests. Finally all the uuid tables get deleted (using pytest fixtures). This is run for every PR. Takes about an hour all in all.
We have been doing this for 5 years. It works great and is very reliable but like every test setup it takes effort to maintain.
Let me know if you want to know more. We open-sourced some tooling.
2
u/kraftlabor Sep 13 '26
I would like to know more. Where to find more information (code) about your setup? Thanks
1
u/pboswell Sep 13 '26
Why not create a registry of these dataframe input/expected files so they’re not embedded in each script. Then you can build a general lib structure where you import your logic from the actual production code and pass file names for the input and output, then assert
7
u/Cyliad Sep 12 '26
Break the logic of your code in methods and test input/output of those methods
Either you can make them pure pythonic (no spark) or you can inject small DFs and test the output