test: core 검색 과 질문 RAG 테스트

pswaao88 · pswaao88 · commit fdab60ad0988 · 2026-06-02T01:28:40.000+09:00
diff --git a/ai/tests/test_core_client.py b/ai/tests/test_core_client.py
@@ -7,6 +7,7 @@
     CoreEmbeddingUpsertError,
     CoreTokenError,
     EmbeddingChunkPayload,
+    EmbeddingSearchHit,
     HttpCoreClient,
 )
 
@@ -241,3 +242,86 @@ async def test_upsert_embeddings_httpx_error_retriable() -> None:
         )
     assert exc_info.value.code == "CORE_UNAVAILABLE"
     assert exc_info.value.retriable is True
+
+
+# ----------- search_embeddings -----------
+
+
+def _make_post_client(
+    *,
+    status: int = 200,
+    json_body: dict | None = None,
+    text: str = "",
+    raise_exc: Exception | None = None,
+) -> MagicMock:
+    client = MagicMock()
+    resp = MagicMock(spec=httpx.Response)
+    resp.status_code = status
+    resp.text = text
+    resp.json = (
+        MagicMock(return_value=json_body)
+        if json_body is not None
+        else MagicMock(side_effect=ValueError("no json"))
+    )
+    if raise_exc is not None:
+        client.post = AsyncMock(side_effect=raise_exc)
+    else:
+        client.post = AsyncMock(return_value=resp)
+    return client
+
+
+@pytest.mark.asyncio
+async def test_search_embeddings_uses_latest_core_contract() -> None:
+    client = _make_post_client(
+        json_body={
+            "hits": [
+                {
+                    "documentId": 7,
+                    "chunkIndex": 2,
+                    "chunkText": "Spring transaction boundaries",
+                    "distance": 0.17,
+                }
+            ]
+        }
+    )
+    core = HttpCoreClient(base_url="http://core:38010", api_key="k", client=client)
+
+    hits = await core.search_embeddings(
+        query_embedding=[0.1, 0.2],
+        document_ids=[7, 8],
+        top_k=3,
+    )
+
+    assert hits == [
+        EmbeddingSearchHit(
+            document_id=7,
+            chunk_index=2,
+            chunk_text="Spring transaction boundaries",
+            distance=0.17,
+        )
+    ]
+    client.post.assert_awaited_once_with(
+        "/api/internal/embeddings/search",
+        json={
+            "queryEmbedding": [0.1, 0.2],
+            "documentIds": [7, 8],
+            "topK": 3,
+        },
+    )
+
+
+@pytest.mark.parametrize("status", [400, 401, 403, 404, 500])
+@pytest.mark.asyncio
+async def test_search_embeddings_non_2xx_returns_empty(status: int) -> None:
+    client = _make_post_client(status=status, text="bad")
+    core = HttpCoreClient(base_url="http://core:38010", api_key="k", client=client)
+
+    assert await core.search_embeddings(query_embedding=[0.1]) == []
+
+
+@pytest.mark.asyncio
+async def test_search_embeddings_http_error_returns_empty() -> None:
+    client = _make_post_client(raise_exc=httpx.ConnectError("dns fail"))
+    core = HttpCoreClient(base_url="http://core:38010", api_key="k", client=client)
+
+    assert await core.search_embeddings(query_embedding=[0.1]) == []
diff --git a/ai/tests/test_followup_consumer.py b/ai/tests/test_followup_consumer.py
@@ -6,6 +6,7 @@
 import pytest
 
 from ai_server.chain.followup_generation_chain import FollowupResult
+from ai_server.core.client import EmbeddingSearchHit
 from ai_server.messaging.consumers.followup_consumer import FollowupConsumer
 from ai_server.messaging.idempotency import LruIdempotencyStore
 from ai_server.model.messages.followup import (
@@ -47,6 +48,7 @@ def _envelope() -> bytes:
             "answerText": "RabbitMQ로 보냈습니다.",
             "mode": "TECHNICAL",
             "jobCategory": "BACKEND",
+            "contextDocumentIds": [7],
         },
         "context": {"userId": 42, "sessionId": 99},
     }
@@ -86,6 +88,85 @@ async def test_consumer_generates_followup_and_publishes_callback():
     assert publisher.publish.await_args.kwargs["message_type"] == "callback.questions"
 
 
+@pytest.mark.asyncio
+async def test_consumer_injects_followup_rag_context_when_available():
+    generator = MagicMock()
+    generator.generate = AsyncMock(
+        return_value=FollowupResult(
+            followup_question="outbox 저장과 발행의 원자성은 어떻게 보장했나요?",
+            answer_evaluation=AnswerEvaluation(
+                specificity=2.0, logic=3.0, structure="PARTIAL_STAR"
+            ),
+        )
+    )
+    publisher = MagicMock()
+    publisher.publish = AsyncMock()
+    core = MagicMock()
+    core.search_embeddings = AsyncMock(
+        return_value=[
+            EmbeddingSearchHit(
+                document_id=7,
+                chunk_index=2,
+                chunk_text="Outbox rows are inserted in the same transaction",
+                distance=0.12,
+            )
+        ]
+    )
+    embedder = MagicMock()
+    embedder.embed = AsyncMock(return_value=[[0.1, 0.2, 0.3]])
+
+    consumer = FollowupConsumer(
+        generator=generator,
+        publisher=publisher,
+        idempotency=LruIdempotencyStore(max_size=10),
+        callback_routing_key="callback.questions",
+        core_client=core,
+        embedder=embedder,
+        rag_top_k=3,
+    )
+    await consumer.handle(_StubMessage(_envelope()))
+
+    embedder.embed.assert_awaited_once()
+    core.search_embeddings.assert_awaited_once_with(
+        query_embedding=[0.1, 0.2, 0.3],
+        document_ids=[7],
+        top_k=3,
+    )
+    context = generator.generate.await_args.kwargs["context"]
+    assert "Outbox rows are inserted in the same transaction" in context
+
+
+@pytest.mark.asyncio
+async def test_consumer_falls_back_when_followup_rag_fails():
+    generator = MagicMock()
+    generator.generate = AsyncMock(
+        return_value=FollowupResult(
+            followup_question="실패 재처리는 어떻게 했나요?",
+            answer_evaluation=AnswerEvaluation(
+                specificity=2.0, logic=3.0, structure="PARTIAL_STAR"
+            ),
+        )
+    )
+    publisher = MagicMock()
+    publisher.publish = AsyncMock()
+    core = MagicMock()
+    core.search_embeddings = AsyncMock(side_effect=RuntimeError("core down"))
+    embedder = MagicMock()
+    embedder.embed = AsyncMock(return_value=[[0.1, 0.2, 0.3]])
+
+    consumer = FollowupConsumer(
+        generator=generator,
+        publisher=publisher,
+        idempotency=LruIdempotencyStore(max_size=10),
+        callback_routing_key="callback.questions",
+        core_client=core,
+        embedder=embedder,
+    )
+    await consumer.handle(_StubMessage(_envelope()))
+
+    assert generator.generate.await_args.kwargs["context"] == "(none)"
+
+
 @pytest.mark.asyncio
 async def test_consumer_idempotent_skip():
     generator = MagicMock()
diff --git a/ai/tests/test_questions_consumer.py b/ai/tests/test_questions_consumer.py
@@ -6,6 +6,7 @@
 import pytest
 
 from ai_server.chain.question_generation_chain import GeneratedQuestionPool
+from ai_server.core.client import EmbeddingSearchHit
 from ai_server.messaging.consumers.questions_consumer import (
     QuestionsConsumer,
     _build_context,
@@ -85,6 +86,7 @@ async def test_consumer_generates_questions_and_publishes_callback():
                 }
             ],
             "maxQuestions": 5,
+            "initialQuestionCount": 2,
         }
     )
     await consumer.handle(_StubMessage(body))
@@ -93,8 +95,8 @@ async def test_consumer_generates_questions_and_publishes_callback():
     call = generator.generate.await_args
     assert call.kwargs["job_category"] == "BACKEND"
     assert call.kwargs["mode"] == "TECHNICAL"
-    # envelope.max_questions(=5) 무시하고 initial_pool_size(default 1) 로 강제. Core 가 첫 질문만 사용.
-    assert call.kwargs["max_questions"] == 1
+    # maxQuestions is the session limit; initialQuestionCount controls this result.
+    assert call.kwargs["max_questions"] == 2
     assert "Java" in call.kwargs["context"]
 
     publisher.publish.assert_awaited_once()
@@ -137,6 +139,166 @@ async def test_consumer_skips_when_message_id_already_seen():
     publisher.publish.assert_not_awaited()
 
 
+@pytest.mark.asyncio
+async def test_consumer_defaults_initial_question_count_to_one():
+    generator = MagicMock()
+    generator.generate = AsyncMock(return_value=GeneratedQuestionPool(questions=[]))
+    publisher = MagicMock()
+    publisher.publish = AsyncMock()
+
+    consumer = QuestionsConsumer(
+        generator=generator,
+        publisher=publisher,
+        idempotency=LruIdempotencyStore(max_size=10),
+        callback_routing_key="callback.questions",
+        initial_pool_size=3,
+    )
+    body = _envelope(
+        {
+            "sessionId": 99,
+            "mode": "TECHNICAL",
+            "jobCategory": "BACKEND",
+            "documents": [],
+            "maxQuestions": 5,
+        }
+    )
+    await consumer.handle(_StubMessage(body))
+
+    assert generator.generate.await_args.kwargs["max_questions"] == 1
+
+
+@pytest.mark.asyncio
+async def test_consumer_clamps_initial_question_count_to_at_least_one():
+    generator = MagicMock()
+    generator.generate = AsyncMock(return_value=GeneratedQuestionPool(questions=[]))
+    publisher = MagicMock()
+    publisher.publish = AsyncMock()
+
+    consumer = QuestionsConsumer(
+        generator=generator,
+        publisher=publisher,
+        idempotency=LruIdempotencyStore(max_size=10),
+        callback_routing_key="callback.questions",
+        initial_pool_size=3,
+    )
+    body = _envelope(
+        {
+            "sessionId": 99,
+            "mode": "TECHNICAL",
+            "jobCategory": "BACKEND",
+            "documents": [],
+            "initialQuestionCount": 0,
+            "maxQuestions": 5,
+        }
+    )
+    await consumer.handle(_StubMessage(body))
+
+    assert generator.generate.await_args.kwargs["max_questions"] == 1
+
+
+@pytest.mark.asyncio
+async def test_consumer_injects_initial_rag_chunks_when_available():
+    generator = MagicMock()
+    generator.generate = AsyncMock(return_value=GeneratedQuestionPool(questions=[]))
+    publisher = MagicMock()
+    publisher.publish = AsyncMock()
+    core = MagicMock()
+    core.search_embeddings = AsyncMock(
+        return_value=[
+            EmbeddingSearchHit(
+                document_id=1,
+                chunk_index=4,
+                chunk_text="Outbox table uses status and retry count",
+                distance=0.11,
+            )
+        ]
+    )
+    embedder = MagicMock()
+    embedder.embed = AsyncMock(return_value=[[0.1, 0.2]])
+
+    consumer = QuestionsConsumer(
+        generator=generator,
+        publisher=publisher,
+        idempotency=LruIdempotencyStore(max_size=10),
+        callback_routing_key="callback.questions",
+        core_client=core,
+        embedder=embedder,
+        rag_top_k=2,
+    )
+    body = _envelope(
+        {
+            "sessionId": 99,
+            "mode": "TECHNICAL",
+            "jobCategory": "BACKEND",
+            "documents": [
+                {
+                    "documentId": 1,
+                    "sourceType": "REPOSITORY",
+                    "summary": "outbox 구현",
+                    "techStack": ["Spring"],
+                    "markdown": "transactional publisher",
+                }
+            ],
+            "initialQuestionCount": 1,
+            "maxQuestions": 5,
+        }
+    )
+    await consumer.handle(_StubMessage(body))
+
+    embedder.embed.assert_awaited_once()
+    core.search_embeddings.assert_awaited_once_with(
+        query_embedding=[0.1, 0.2],
+        document_ids=[1],
+        top_k=2,
+    )
+    context = generator.generate.await_args.kwargs["context"]
+    assert "Outbox table uses status and retry count" in context
+    assert "outbox 구현" in context
+
+
+@pytest.mark.asyncio
+async def test_consumer_falls_back_to_document_context_when_rag_fails():
+    generator = MagicMock()
+    generator.generate = AsyncMock(return_value=GeneratedQuestionPool(questions=[]))
+    publisher = MagicMock()
+    publisher.publish = AsyncMock()
+    core = MagicMock()
+    core.search_embeddings = AsyncMock(side_effect=RuntimeError("core down"))
+    embedder = MagicMock()
+    embedder.embed = AsyncMock(return_value=[[0.1, 0.2]])
+
+    consumer = QuestionsConsumer(
+        generator=generator,
+        publisher=publisher,
+        idempotency=LruIdempotencyStore(max_size=10),
+        callback_routing_key="callback.questions",
+        core_client=core,
+        embedder=embedder,
+    )
+    body = _envelope(
+        {
+            "sessionId": 99,
+            "mode": "TECHNICAL",
+            "jobCategory": "BACKEND",
+            "documents": [
+                {
+                    "documentId": 1,
+                    "sourceType": "RESUME",
+                    "summary": "Java/Spring backend",
+                    "techStack": ["Java"],
+                    "markdown": "payment service",
+                }
+            ],
+            "maxQuestions": 5,
+        }
+    )
+    await consumer.handle(_StubMessage(body))
+
+    context = generator.generate.await_args.kwargs["context"]
+    assert "Java/Spring backend" in context
+    assert "Retrieved document chunks" not in context
+
+
 def test_build_context_handles_empty_documents():
     assert _build_context([]) == "(no documents)"
 
