diff --git a/exercises/solutions/04_unit_tests_solutions.ipynb b/exercises/solutions/04_unit_tests_solutions.ipynb
index 8e1e9a2..93bc1b4 100644
--- a/exercises/solutions/04_unit_tests_solutions.ipynb
+++ b/exercises/solutions/04_unit_tests_solutions.ipynb
@@ -11,9 +11,26 @@
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": null,
    "id": "1",
    "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Setup code for solutions\n",
+    "\n",
+    "import os\n",
+    "import sys\n",
+    "\n",
+    "import pandas as pd\n",
+    "\n",
+    "sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), \"..\", \"..\", \"src\")))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2",
+   "metadata": {},
    "source": [
     "## Exercise 1 solution: Review and adapt an existing unit test\n",
     "\n",
@@ -27,7 +44,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2",
+   "id": "3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -57,47 +74,29 @@
   },
   {
    "cell_type": "markdown",
-   "id": "3",
+   "id": "4",
    "metadata": {},
    "source": [
     "Notice how the check for the missing smoker value fails. The test checks the first column for a \"smoker\" value of \"No\", however the modified `clean_health_data` function fills missing smoker values with 'Yes', changing the smoker value in the first column to 'Yes', which causes the test to fail.\n",
     "\n",
     "The test failing highlights the function to developers who can then check if the change was correct or not. If it was not, the developer can fix the error in the function. If it was, the unit test can be adapted to incorporate the change. In this case assume the change was correct. In order for the test to pass change the expected \"smoker\" value to \"Yes\" instead of \"No\".\n",
-    "The original assert statement looks like this:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "4",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "assert cleaned[\"smoker\"].iloc[0] == \"No\""
+    "The original assert statement looks like this:  \n",
+    "  \n",
+    "```python\n",
+    "assert cleaned[\"smoker\"].iloc[0] == \"No\"\n",
+    "```  \n",
+    "  \n",
+    "The changed assert statement should look like this:  \n",
+    "  \n",
+    "```python\n",
+    "assert cleaned[\"smoker\"].iloc[0] == \"Yes\"\n",
+    "```"
    ]
   },
   {
    "cell_type": "markdown",
    "id": "5",
    "metadata": {},
-   "source": [
-    "The changed assert statement should look like this:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "6",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "assert cleaned[\"smoker\"].iloc[0] == \"Yes\""
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "7",
-   "metadata": {},
    "source": [
     "## Exercise 2 Solution: Write a simple unit test for a new function\n",
     "\n",
@@ -107,21 +106,16 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8",
+   "id": "6",
    "metadata": {},
    "outputs": [],
    "source": [
     "# Walkthrough: Unit test for flag_missing\n",
-    "import os\n",
-    "import sys\n",
-    "\n",
-    "import pandas as pd\n",
     "\n",
-    "sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), \"..\", \"..\", \"src\")))\n",
+    "# Note: This test will not run unless flag_missing has been entered into cleaning.py\n",
     "\n",
     "from python_rap_demo.cleaning import flag_missing\n",
     "\n",
-    "\n",
     "def test_flag_missing():\n",
     "    \"\"\"\n",
     "    Test flag_missing\n",
@@ -141,7 +135,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9",
+   "id": "7",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -149,14 +143,6 @@
     "\n",
     "# Note: This test will not run unless impute_by_group has been entered into cleaning.py\n",
     "\n",
-    "\n",
-    "import os\n",
-    "import sys\n",
-    "\n",
-    "import pandas as pd\n",
-    "\n",
-    "sys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), \"..\", \"..\", \"src\")))\n",
-    "\n",
     "from python_rap_demo.cleaning import impute_by_group\n",
     "\n",
     "\n",
@@ -177,7 +163,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "10",
+   "id": "8",
    "metadata": {},
    "source": [
     "## Exercise 3 solution: Run your unit tests\n",
@@ -192,7 +178,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "11",
+   "id": "9",
    "metadata": {},
    "source": [
     "## Exercise 4 solution: Stretch - Check test coverage\n",
@@ -209,7 +195,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "12",
+   "id": "10",
    "metadata": {},
    "source": [
     "## Exercise 5 solution: Stretch - Try parameterisation in pytest\n",
@@ -220,7 +206,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "13",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -311,7 +297,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.3"
+   "version": "3.12.5"
   }
  },
  "nbformat": 4,