ml

spotify_silver_gold.ipynb

@@ -363,6 +363,88 @@
     ").orderBy(F.desc(\"skip_prob\")).limit(10))"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6cd50786",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## --- Tests for ETL Functions --- ##\n",
+    "import pytest\n",
+    "\n",
+    "# Helper: collect a single column into a Python list\n",
+    "def col_vals(df, colname):\n",
+    "    return [row[colname] for row in df.select(colname).collect()]\n",
+    "\n",
+    "# ── Test 1: clean_platform ──\n",
+    "def test_clean_platform():\n",
+    "    \"\"\"Platform strings are normalized into 5 buckets.\"\"\"\n",
+    "    sample = spark.createDataFrame([\n",
+    "        (\"iOS 14.2 (iPhone)\",),\n",
+    "        (\"Android OS 11 API 30\",),\n",
+    "        (\"Windows 10\",),\n",
+    "        (\"macOS 12.0\",),\n",
+    "        (\"web player\",),\n",
+    "        (\"Partner amazon_salmon\",),\n",
+    "    ], [\"platform\"])\n",
+    "    result = clean_platform(sample)\n",
+    "    vals = col_vals(result, \"platform\")\n",
+    "    assert vals == [\"ios\", \"android\", \"desktop\", \"desktop\", \"web\", \"other\"], f\"Got {vals}\"\n",
+    "    print(\"[x] test_clean_platform passed\")\n",
+    "\n",
+    "# ── Test 2: derive_skip_label ──\n",
+    "def test_derive_skip_label():\n",
+    "    \"\"\"Skip = fwdbtn/backbtn with ms_played < 30s; everything else is not a skip.\"\"\"\n",
+    "    sample = spark.createDataFrame([\n",
+    "        (\"fwdbtn\",    500.0),\n",
+    "        (\"backbtn\",   200.0),\n",
+    "        (\"trackdone\", 180000.0),\n",
+    "        (\"fwdbtn\",    60000.0),   # fwdbtn but >= 30s → not a skip\n",
+    "        (\"endplay\",   100.0),   # wrong reason → not a skip\n",
+    "    ], [\"reason_end\", \"ms_played\"])\n",
+    "    result = derive_skip_label(sample)\n",
+    "    vals = col_vals(result, \"is_skip\")\n",
+    "    assert vals == [1, 1, 0, 0, 0], f\"Got {vals}\"\n",
+    "    print(\"[x] test_derive_skip_label passed\")\n",
+    "\n",
+    "# ── Test 3: extract_track_id ──\n",
+    "def test_extract_track_id():\n",
+    "    \"\"\"Track ID is extracted from the spotify:track: URI.\"\"\"\n",
+    "    sample = spark.createDataFrame([\n",
+    "        (\"spotify:track:0cKk8BKEi7zXbdrYdyqBP5\",),\n",
+    "        (\"spotify:track:2VgPA0WnwvE2nbltxa3BHz\",),\n",
+    "        (\"\",),   # empty URI → empty string\n",
+    "    ], [\"track_uri\"])\n",
+    "    result = extract_track_id(sample)\n",
+    "    vals = col_vals(result, \"track_id\")\n",
+    "    assert vals[0] == \"0cKk8BKEi7zXbdrYdyqBP5\"\n",
+    "    assert vals[1] == \"2VgPA0WnwvE2nbltxa3BHz\"\n",
+    "    assert vals[2] == \"\"\n",
+    "    print(\"[x] test_extract_track_id passed\")\n",
+    "\n",
+    "# ── Test 4: add_time_features ──\n",
+    "def test_add_time_features():\n",
+    "    \"\"\"Hour and day-of-week are derived from the ts timestamp.\"\"\"\n",
+    "    sample = spark.createDataFrame([\n",
+    "        (\"2021-10-27 20:17:11\",),  # Wed, 8pm\n",
+    "        (\"2022-01-03 07:30:00\",),  # Mon, 7am\n",
+    "    ], [\"ts\"])\n",
+    "    result = add_time_features(sample)\n",
+    "    hours = col_vals(result, \"hour\")\n",
+    "    dows  = col_vals(result, \"dow\")\n",
+    "    assert hours == [20, 7], f\"Got hours {hours}\"\n",
+    "    assert dows  == [4, 2], f\"Got dows {dows}\"   # Sun=1, Mon=2, …, Wed=4\n",
+    "    print(\"[x] test_add_time_features passed\")\n",
+    "\n",
+    "# ── Run all tests ──\n",
+    "test_clean_platform()\n",
+    "test_derive_skip_label()\n",
+    "test_extract_track_id()\n",
+    "test_add_time_features()\n",
+    "print(\"\\n🎉 All 4 ETL tests passed!\")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "23c744a6",