From c05296376eb2fa8607d9f1f61039c7c9e21ea73a Mon Sep 17 00:00:00 2001
From: cybermaggedon <cybermaggedon@gmail.com>
Date: Fri, 3 Jul 2026 13:44:09 +0100
Subject: [PATCH] fix: remove test import (#1017)

Convention in the tests is to just import the libraries as production
code would. This is fragile, and could possibly be used to inject
malicious code in the CI environment.
---
 .../test_document_rag_diversity_selection.py  | 28 ++-----------------
 1 file changed, 3 insertions(+), 25 deletions(-)

diff --git a/tests/unit/test_retrieval/test_document_rag_diversity_selection.py b/tests/unit/test_retrieval/test_document_rag_diversity_selection.py
index 47d03cb8..6dcd9458 100644
--- a/tests/unit/test_retrieval/test_document_rag_diversity_selection.py
+++ b/tests/unit/test_retrieval/test_document_rag_diversity_selection.py
@@ -1,30 +1,8 @@
-import importlib.util
-from pathlib import Path
-
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-RERANK_PATH = (
-    REPO_ROOT
-    / "trustgraph-flow"
-    / "trustgraph"
-    / "retrieval"
-    / "document_rag"
-    / "rerank.py"
+from trustgraph.retrieval.document_rag.rerank import (
+    RerankCandidate, normalize_candidate_scores, mmr_select,
+    _pair_diversity_penalty
 )
 
-spec = importlib.util.spec_from_file_location(
-    "document_rag_diversity_rerank",
-    RERANK_PATH,
-)
-rerank = importlib.util.module_from_spec(spec)
-spec.loader.exec_module(rerank)
-
-RerankCandidate = rerank.RerankCandidate
-normalize_candidate_scores = rerank.normalize_candidate_scores
-mmr_select = rerank.mmr_select
-_pair_diversity_penalty = rerank._pair_diversity_penalty
-
-
 def candidate(index, chunk_id, text, score):
     return RerankCandidate(
         index=index,