diff --git a/dp_wizard/app/components/column_module.py b/dp_wizard/app/components/column_module.py
index 47054b9..d46fe9c 100644
--- a/dp_wizard/app/components/column_module.py
+++ b/dp_wizard/app/components/column_module.py
@@ -186,7 +186,10 @@ def column_plot():
             weighted_epsilon=epsilon * weight / weights_sum,
         )
         s = "s" if contributions > 1 else ""
-        title = f"Simulated {name}: normal distribution, {contributions} contribution{s} / invidual"
+        title = (
+            f"Simulated {name}: normal distribution, "
+            f"{contributions} contribution{s} / invidual"
+        )
         return plot_histogram(
             histogram,
             error=accuracy,
diff --git a/dp_wizard/utils/code_generators/no-tests/_privacy_unit.py b/dp_wizard/utils/code_generators/no-tests/_privacy_unit.py
index a15e480..7e43432 100644
--- a/dp_wizard/utils/code_generators/no-tests/_privacy_unit.py
+++ b/dp_wizard/utils/code_generators/no-tests/_privacy_unit.py
@@ -1 +1,2 @@
-privacy_unit = dp.unit_of(contributions=CONTRIBUTIONS)
+contributions = CONTRIBUTIONS
+privacy_unit = dp.unit_of(contributions=contributions)
diff --git a/tests/test_app.py b/tests/test_app.py
index f4e4c40..7aa2209 100644
--- a/tests/test_app.py
+++ b/tests/test_app.py
@@ -54,7 +54,7 @@ def expect_no_error():
     expect_not_visible(download_results_text)
     page.get_by_label("Contributions").fill("42")
     page.get_by_text("Code sample: Unit of Privacy").click()
-    expect_visible("dp.unit_of(contributions=42)")
+    expect_visible("contributions = 42")
     expect_no_error()
 
     # Button disabled until upload:
@@ -156,7 +156,7 @@ def expect_no_error():
 
     script_download = script_download_info.value
     script = script_download.path().read_text()
-    assert "privacy_unit = dp.unit_of(contributions=42)" in script
+    assert "contributions = 42" in script
 
     # Notebook:
     with page.expect_download() as notebook_download_info:
@@ -165,7 +165,7 @@ def expect_no_error():
 
     notebook_download = notebook_download_info.value
     notebook = notebook_download.path().read_text()
-    assert "privacy_unit = dp.unit_of(contributions=42)" in notebook
+    assert "contributions = 42" in notebook
 
     # -- Feedback --
     page.get_by_text("Feedback").click()