wehub-resource-sync
6b7e6b44f1
gh-pages / build (push) Waiting to run
Python Publish (pypi) / Upload release to PyPI (push) Waiting to run
Spellcheck / spellcheck (push) Waiting to run
Python Build and Type Check / python-ci (ubuntu-latest, 3.11) (push) Has been cancelled
Python Build and Type Check / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Build and Type Check / python-ci (windows-latest, 3.11) (push) Has been cancelled
Python Build and Type Check / python-ci (windows-latest, 3.13) (push) Has been cancelled
Python Integration Tests / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Integration Tests / python-ci (windows-latest, 3.13) (push) Has been cancelled
Python Notebook Tests / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Notebook Tests / python-ci (windows-latest, 3.13) (push) Has been cancelled
Python Smoke Tests / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Smoke Tests / python-ci (windows-latest, 3.13) (push) Has been cancelled
Python Unit Tests / python-ci (ubuntu-latest, 3.13) (push) Has been cancelled
Python Unit Tests / python-ci (windows-latest, 3.13) (push) Has been cancelled
66 lines
1.9 KiB
Python
66 lines
1.9 KiB
Python
# Copyright (C) 2026 Microsoft
|
|
# Licensed under the MIT License
|
|
|
|
"""Verb test for the update_text_embeddings workflow."""
|
|
|
|
from unittest.mock import patch
|
|
|
|
from graphrag.config.embeddings import all_embeddings
|
|
from graphrag.index.workflows.update_text_embeddings import (
|
|
run_workflow,
|
|
)
|
|
|
|
from tests.unit.config.utils import get_default_graphrag_config
|
|
|
|
from .util import create_test_context
|
|
|
|
|
|
async def test_update_text_embeddings():
|
|
"""Verify update_text_embeddings produces embedding tables.
|
|
|
|
Mocks get_update_table_providers to return the test context's
|
|
output_table_provider, simulating the merged tables written by
|
|
upstream update workflows.
|
|
"""
|
|
context = await create_test_context(
|
|
storage=[
|
|
"documents",
|
|
"relationships",
|
|
"text_units",
|
|
"entities",
|
|
"community_reports",
|
|
]
|
|
)
|
|
context.state["update_timestamp"] = "20260220-000000"
|
|
|
|
config = get_default_graphrag_config()
|
|
llm_settings = config.get_embedding_model_config(
|
|
config.embed_text.embedding_model_id
|
|
)
|
|
llm_settings.type = "mock"
|
|
llm_settings.mock_responses = [1.0] * 3072
|
|
|
|
config.embed_text.names = list(all_embeddings)
|
|
config.snapshots.embeddings = True
|
|
|
|
with patch(
|
|
"graphrag.index.workflows.update_text_embeddings.get_update_table_providers",
|
|
) as mock_providers:
|
|
mock_providers.return_value = (
|
|
context.output_table_provider,
|
|
None,
|
|
None,
|
|
)
|
|
await run_workflow(config, context)
|
|
|
|
parquet_files = context.output_storage.keys()
|
|
for field in all_embeddings:
|
|
assert f"embeddings.{field}.parquet" in parquet_files
|
|
|
|
entity_embeddings = await context.output_table_provider.read_dataframe(
|
|
"embeddings.entity_description"
|
|
)
|
|
assert len(entity_embeddings.columns) == 2
|
|
assert "id" in entity_embeddings.columns
|
|
assert "embedding" in entity_embeddings.columns
|