From d432276ee845719f0c2703f400334bec9445be12 Mon Sep 17 00:00:00 2001
From: Fabrizio Waldner <fabrizio@nebius.com>
Date: Thu, 16 Jul 2026 17:19:43 +0200
Subject: [PATCH] Run Python unit tests in CI

---
 .github/workflows/one_job.yml     | 15 +++++++++++++++
 Makefile                          | 24 ++++++++++++++++++++++--
 images/worker/worker_init_test.py |  8 ++++++--
 3 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/one_job.yml b/.github/workflows/one_job.yml
index 6d95076cb..d6062c552 100644
--- a/.github/workflows/one_job.yml
+++ b/.github/workflows/one_job.yml
@@ -204,6 +204,20 @@ jobs:
           fi
           go build "${packages[@]}"
 
+  test-python:
+    needs: [changes, authorize]
+    if: needs.changes.outputs.should_build == 'true'
+    runs-on: ubuntu-24.04
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
+
+      - name: Test Python
+        run: make test-python
+
   build-soperator-images:
     needs:
       - changes
@@ -846,6 +860,7 @@ jobs:
       - pre-build
       - lint
       - build-e2e
+      - test-python
       - build-slurm-images
       - build-soperator-images
       - manifest-populate-jail
diff --git a/Makefile b/Makefile
index af189e2f4..900079306 100644
--- a/Makefile
+++ b/Makefile
@@ -120,8 +120,28 @@ test: manifests generate fmt vet envtest ## Run tests.
 	go test ./...
 
 .PHONY: test-python
-test-python: ## Temporarily retained while removing Python unit tests from CI.
-	@echo "Python unit tests are disabled in CI."
+test-python: ## Run all Python unit tests.
+	@found=false; \
+	while IFS= read -r -d '' test_file; do \
+		found=true; \
+		test_dir=$$(dirname "$$test_file"); \
+		test_pattern=$$(basename "$$test_file"); \
+		echo "Running Python unit tests in $$test_file"; \
+		PYTHONDONTWRITEBYTECODE=1 python3 -m unittest discover \
+			-s "$$test_dir" \
+			-p "$$test_pattern" \
+			-v; \
+	done < <(find . \
+		-type d \( \
+			-name .git -o \
+			-name .venv -o \
+			-name venv -o \
+			-name __pycache__ \
+		\) -prune -o \
+		-type f \( -name '*_test.py' -o -name 'test_*.py' \) -print0); \
+	if [ "$$found" = false ]; then \
+		echo "No Python unit tests found; skipping."; \
+	fi
 
 .PHONY: test-coverage
 test-coverage: manifests generate fmt vet envtest ## Run tests and generate test coverage.
diff --git a/images/worker/worker_init_test.py b/images/worker/worker_init_test.py
index ae71d4a51..22ab88c12 100644
--- a/images/worker/worker_init_test.py
+++ b/images/worker/worker_init_test.py
@@ -895,7 +895,9 @@ def test_non_gpu_applies_per_fabric_unknown_topology(
             worker_init.wait_for_topology()
 
         mock_apply.assert_called_once_with(
-            "worker-0", "topology=default:fab-a:fab-a.unknown"
+            "worker-0",
+            "topology=default:fab-a:fab-a.unknown",
+            worker_init.TOPOLOGY_PLUGIN_TREE,
         )
 
     @mock.patch("worker_init.wait_for_hostname_in_topology_conf")
@@ -915,7 +917,9 @@ def test_non_gpu_applies_per_fabric_unknown_block(
             worker_init.wait_for_topology()
 
         mock_apply.assert_called_once_with(
-            "worker-0", "topology=default:fab-a.unknown"
+            "worker-0",
+            "topology=default:fab-a.unknown",
+            worker_init.TOPOLOGY_PLUGIN_BLOCK,
         )
 
     @mock.patch("worker_init.wait_for_hostname_in_topology_conf")