From 4519069f5d8e69e62b84411f0be84a41508ea2f8 Mon Sep 17 00:00:00 2001 From: Codex Date: Mon, 21 Sep 2026 05:07:09 +0200 Subject: [PATCH] Batch neighborhood traversal queries --- app/services/material/neighbor_service.py | 189 +++++++++++++++++- app/services/material/neighborhood_service.py | 141 +++++++------ .../neighborhood_api_remediation.md | 79 ++++++++ ...eighborhood_api_remediation_benchmark.json | 56 ++++++ scripts/diagnose_neighborhood_performance.py | 17 +- .../material/test_neighbor_batch_service.py | 143 +++++++++++++ .../material/test_neighborhood_service.py | 42 ++-- 7 files changed, 570 insertions(+), 97 deletions(-) create mode 100644 docs/performance/neighborhood_api_remediation.md create mode 100644 docs/performance/neighborhood_api_remediation_benchmark.json create mode 100644 tests/services/material/test_neighbor_batch_service.py diff --git a/app/services/material/neighbor_service.py b/app/services/material/neighbor_service.py index 451668a..74a58dc 100644 --- a/app/services/material/neighbor_service.py +++ b/app/services/material/neighbor_service.py @@ -29,7 +29,6 @@ def get_neighbors(self, material_id: int) -> dict: element_ids = self._get_material_element_ids(material_id) application_ids = self._get_material_application_ids(material_id) - neighbor_scores: dict[int, dict] = {} self._collect_element_neighbors( @@ -37,7 +36,6 @@ def get_neighbors(self, material_id: int) -> dict: element_ids=element_ids, neighbor_scores=neighbor_scores, ) - self._collect_application_neighbors( material_id=material_id, application_ids=application_ids, @@ -45,16 +43,87 @@ def get_neighbors(self, material_id: int) -> dict: ) materials_by_id = self._get_materials_by_id( - material_ids=list(neighbor_scores.keys()) + material_ids=list(neighbor_scores) ) - neighbors = self._build_neighbors( neighbor_scores=neighbor_scores, materials_by_id=materials_by_id, ) - neighbors.sort(key=neighbor_ranking_key) + return self._neighbors_response( + material=material, + neighbors=neighbors, + ) + + def get_neighbors_batch(self, material_ids: list[int]) -> dict[int, dict]: + ordered_material_ids = list(dict.fromkeys(material_ids)) + + if not ordered_material_ids: + return {} + + materials_by_id = { + material.id: material + for material in self.db.query(Material) + .filter(Material.id.in_(ordered_material_ids)) + .all() + } + responses = { + material_id: self._empty_neighbors_response(material_id) + for material_id in ordered_material_ids + } + + if not materials_by_id: + return responses + + source_material_ids = set(materials_by_id) + element_ids_by_material = self._get_element_ids_by_material( + source_material_ids + ) + application_ids_by_material = self._get_application_ids_by_material( + source_material_ids + ) + neighbor_scores_by_material = { + material_id: {} + for material_id in source_material_ids + } + + self._collect_element_neighbors_batch( + element_ids_by_material=element_ids_by_material, + neighbor_scores_by_material=neighbor_scores_by_material, + ) + self._collect_application_neighbors_batch( + application_ids_by_material=application_ids_by_material, + neighbor_scores_by_material=neighbor_scores_by_material, + ) + + neighbor_material_ids = { + neighbor_id + for neighbor_scores in neighbor_scores_by_material.values() + for neighbor_id in neighbor_scores + } + neighbor_materials_by_id = self._get_materials_by_id( + material_ids=list(neighbor_material_ids) + ) + + for material_id, material in materials_by_id.items(): + neighbors = self._build_neighbors( + neighbor_scores=neighbor_scores_by_material[material_id], + materials_by_id=neighbor_materials_by_id, + ) + neighbors.sort(key=neighbor_ranking_key) + responses[material_id] = self._neighbors_response( + material=material, + neighbors=neighbors, + ) + + return responses + + def _neighbors_response( + self, + material: Material, + neighbors: list[dict], + ) -> dict: return { "material_id": material.id, "mp_id": material.mp_id, @@ -66,6 +135,114 @@ def get_neighbors(self, material_id: int) -> dict: "neighbors": neighbors, } + def _get_element_ids_by_material( + self, + material_ids: set[int], + ) -> dict[int, set[int]]: + element_ids_by_material = { + material_id: set() + for material_id in material_ids + } + rows = ( + self.db.query(MaterialElement) + .filter(MaterialElement.material_id.in_(material_ids)) + .all() + ) + + for row in rows: + element_ids_by_material[row.material_id].add(row.element_id) + + return element_ids_by_material + + def _get_application_ids_by_material( + self, + material_ids: set[int], + ) -> dict[int, set[int]]: + application_ids_by_material = { + material_id: set() + for material_id in material_ids + } + rows = ( + self.db.query(MaterialApplication) + .filter(MaterialApplication.material_id.in_(material_ids)) + .all() + ) + + for row in rows: + application_ids_by_material[row.material_id].add(row.application_id) + + return application_ids_by_material + + def _collect_element_neighbors_batch( + self, + element_ids_by_material: dict[int, set[int]], + neighbor_scores_by_material: dict[int, dict[int, dict]], + ) -> None: + source_ids_by_element: dict[int, set[int]] = {} + for material_id, element_ids in element_ids_by_material.items(): + for element_id in element_ids: + source_ids_by_element.setdefault(element_id, set()).add(material_id) + + if not source_ids_by_element: + return + + rows = ( + self.db.query(MaterialElement) + .filter(MaterialElement.element_id.in_(source_ids_by_element)) + .all() + ) + + for row in rows: + for source_material_id in source_ids_by_element[row.element_id]: + if row.material_id == source_material_id: + continue + + score_data = self._get_or_create_score_data( + material_id=row.material_id, + neighbor_scores=neighbor_scores_by_material[source_material_id], + ) + score_data["shared_element_count"] += 1 + score_data["relationship_types"].add(self.ELEMENT_RELATIONSHIP) + + def _collect_application_neighbors_batch( + self, + application_ids_by_material: dict[int, set[int]], + neighbor_scores_by_material: dict[int, dict[int, dict]], + ) -> None: + source_ids_by_application: dict[int, set[int]] = {} + for material_id, application_ids in application_ids_by_material.items(): + for application_id in application_ids: + source_ids_by_application.setdefault(application_id, set()).add( + material_id + ) + + if not source_ids_by_application: + return + + rows = ( + self.db.query(MaterialApplication) + .filter( + MaterialApplication.application_id.in_( + source_ids_by_application + ) + ) + .all() + ) + + for row in rows: + for source_material_id in source_ids_by_application[row.application_id]: + if row.material_id == source_material_id: + continue + + score_data = self._get_or_create_score_data( + material_id=row.material_id, + neighbor_scores=neighbor_scores_by_material[source_material_id], + ) + score_data["shared_application_count"] += 1 + score_data["relationship_types"].add( + self.APPLICATION_RELATIONSHIP + ) + def _empty_neighbors_response(self, material_id: int) -> dict: return { "material_id": material_id, @@ -210,4 +387,4 @@ def _calculate_neighbor_score(self, score_data: dict) -> int: return ( score_data["shared_element_count"] * 2 + score_data["shared_application_count"] * 3 - ) \ No newline at end of file + ) diff --git a/app/services/material/neighborhood_service.py b/app/services/material/neighborhood_service.py index 4b40887..71cc5a8 100644 --- a/app/services/material/neighborhood_service.py +++ b/app/services/material/neighborhood_service.py @@ -1,5 +1,3 @@ -from collections import deque - from sqlalchemy.orm import Session from app.services.material.neighbor_service import ( @@ -21,10 +19,10 @@ def get_neighborhood( ) -> dict: neighbor_cache: dict[int, dict] = {} - root = self._get_neighbors( - material_id=material_id, + root = self._get_neighbors_batch( + material_ids=[material_id], cache=neighbor_cache, - ) + )[material_id] if root["mp_id"] is None: return self._empty_neighborhood_response( @@ -33,7 +31,7 @@ def get_neighborhood( ) visited: set[int] = {material_id} - frontier: deque[tuple[int, int]] = deque([(material_id, 0)]) + frontier: list[int] = [material_id] nodes: dict[int, dict] = { material_id: { @@ -51,64 +49,69 @@ def get_neighborhood( edges: list[dict] = [] - while frontier: - current_id, current_depth = frontier.popleft() - - if current_depth >= depth: - continue - - current_neighbors = self._get_neighbors( - material_id=current_id, + current_depth = 0 + while frontier and current_depth < depth: + neighbors_by_material = self._get_neighbors_batch( + material_ids=frontier, cache=neighbor_cache, ) + next_frontier: list[int] = [] - ordered_neighbors = sorted( - current_neighbors["neighbors"], - key=neighbor_ranking_key, - ) + for current_id in frontier: + current_neighbors = neighbors_by_material[current_id] + + ordered_neighbors = sorted( + current_neighbors["neighbors"], + key=neighbor_ranking_key, + ) - for neighbor in ordered_neighbors: - neighbor_id = neighbor["material_id"] - next_depth = current_depth + 1 - - if neighbor_id not in nodes: - if len(nodes) >= limit: - continue - - nodes[neighbor_id] = { - "material_id": neighbor_id, - "mp_id": neighbor["mp_id"], - "pretty_formula": neighbor["pretty_formula"], - "formula": neighbor["formula"], - "material_type": neighbor["material_type"], - "is_stable": neighbor["is_stable"], - "energy_above_hull": neighbor["energy_above_hull"], - "depth": next_depth, - "best_score": neighbor["neighbor_score"], - } - - else: - nodes[neighbor_id]["best_score"] = max( - nodes[neighbor_id]["best_score"], - neighbor["neighbor_score"], + for neighbor in ordered_neighbors: + neighbor_id = neighbor["material_id"] + next_depth = current_depth + 1 + + if neighbor_id not in nodes: + if len(nodes) >= limit: + continue + + nodes[neighbor_id] = { + "material_id": neighbor_id, + "mp_id": neighbor["mp_id"], + "pretty_formula": neighbor["pretty_formula"], + "formula": neighbor["formula"], + "material_type": neighbor["material_type"], + "is_stable": neighbor["is_stable"], + "energy_above_hull": neighbor["energy_above_hull"], + "depth": next_depth, + "best_score": neighbor["neighbor_score"], + } + + else: + nodes[neighbor_id]["best_score"] = max( + nodes[neighbor_id]["best_score"], + neighbor["neighbor_score"], + ) + + edges.append( + { + "source_material_id": current_id, + "target_material_id": neighbor_id, + "relationship_types": neighbor["relationship_types"], + "shared_element_count": neighbor[ + "shared_element_count" + ], + "shared_application_count": neighbor[ + "shared_application_count" + ], + "edge_score": neighbor["neighbor_score"], + } ) - edges.append( - { - "source_material_id": current_id, - "target_material_id": neighbor_id, - "relationship_types": neighbor["relationship_types"], - "shared_element_count": neighbor["shared_element_count"], - "shared_application_count": neighbor[ - "shared_application_count" - ], - "edge_score": neighbor["neighbor_score"], - } - ) + if neighbor_id not in visited: + visited.add(neighbor_id) + next_frontier.append(neighbor_id) - if neighbor_id not in visited: - visited.add(neighbor_id) - frontier.append((neighbor_id, next_depth)) + frontier = next_frontier + current_depth += 1 sorted_nodes = sorted( nodes.values(), @@ -149,15 +152,25 @@ def get_neighborhood( "edges": limited_edges, } - def _get_neighbors( + def _get_neighbors_batch( self, - material_id: int, + material_ids: list[int], cache: dict[int, dict], - ) -> dict: - if material_id not in cache: - cache[material_id] = self.neighbor_service.get_neighbors(material_id) + ) -> dict[int, dict]: + missing_material_ids = [ + material_id + for material_id in material_ids + if material_id not in cache + ] + if missing_material_ids: + cache.update( + self.neighbor_service.get_neighbors_batch(missing_material_ids) + ) - return cache[material_id] + return { + material_id: cache[material_id] + for material_id in material_ids + } def _empty_neighborhood_response( self, @@ -174,4 +187,4 @@ def _empty_neighborhood_response( "edge_count": 0, "nodes": [], "edges": [], - } \ No newline at end of file + } diff --git a/docs/performance/neighborhood_api_remediation.md b/docs/performance/neighborhood_api_remediation.md new file mode 100644 index 0000000..3644ab2 --- /dev/null +++ b/docs/performance/neighborhood_api_remediation.md @@ -0,0 +1,79 @@ +# Neighborhood API Performance Remediation + +**Status:** Local implementation and acceptance benchmarking complete + +**Scope:** `GET /api/v1/materials/{material_id}/neighborhood` + +**Production access or mutation:** None + +## Outcome + +The neighborhood traversal now preloads neighbor inputs once per BFS level and +then applies the existing sequential admission, scoring, ordering, edge +closure, and response assembly logic. For the production-count synthetic +fixture, the default depth-two request fell from 150 SQL queries to 12 while +producing the same 9,657-byte payload and payload SHA-256. + +| Measurement | Diagnostic baseline | Remediation | Result | +|---|---:|---:|---:| +| SQL queries | 150 | 12 | 92% reduction | +| Warm median service latency | 1,638.338 ms | 342.218 ms | 4.79× faster | +| 100 ms/query sensitivity | 16,749.582 ms | 1,756.175 ms | 9.54× faster | +| Payload | 9,657 bytes | 9,657 bytes | unchanged | +| Payload SHA-256 | `50f22e...fd11` | `50f22e...fd11` | unchanged | + +The committed machine-readable results are in +`docs/performance/neighborhood_api_remediation_benchmark.json`. + +## Implementation + +`MaterialNeighborService.get_neighbors_batch()`: + +1. fetches all requested source materials; +2. fetches their element memberships; +3. fetches their application memberships; +4. fetches matching element associations once for the batch; +5. fetches matching application associations once for the batch; and +6. fetches all resulting neighbor materials once. + +It reconstructs the same per-source score dictionaries and delegates neighbor +construction to the existing `_build_neighbors()` and score formula. + +`MaterialNeighborhoodService` batches only data loading. It still processes +each frontier material and each ranked neighbor in the same deterministic BFS +sequence. Batching therefore does not change which material is admitted when +the node limit is reached. + +## Semantic evidence + +- Batched neighbor results equal the legacy single-material results across + shared-element, shared-application, combined, isolated, null-valued, and + missing-material cases. +- Tied rankings remain resolved by material ID. +- Existing depth, node-limit, deterministic-order, edge-closure, and count + tests pass with the batch-call boundary asserted. +- Four production-count default benchmark runs were identical. +- The default endpoint payload hash exactly matches the diagnostic baseline. +- Payload hashes at limits 1, 5, 10, 25, 50, and 100 are deterministic. + +## Acceptance results + +All locally executable acceptance gates passed: + +- depth-two query count is 12 and does not grow with limits above one; +- default warm median is below 500 ms; +- default response at 100 ms simulated query latency is below 3 seconds; +- serialization remains below 5 ms; +- payload compatibility is unchanged; and +- limits 50 and 100 complete below the 20-second request deadline. + +The limit-100 request completed in 1,720.358 ms warm and 2,823.650 ms with +100 ms simulated latency. + +## Remaining boundary + +Disposable PostgreSQL validation with the exact accepted 1,727-material +manifest remains an operational acceptance step where that environment is +available. No PostgreSQL performance claim, concurrency claim, deployment, +restart, production read, database write, schema change, migration, or index +change is included in this remediation. diff --git a/docs/performance/neighborhood_api_remediation_benchmark.json b/docs/performance/neighborhood_api_remediation_benchmark.json new file mode 100644 index 0000000..d5da623 --- /dev/null +++ b/docs/performance/neighborhood_api_remediation_benchmark.json @@ -0,0 +1,56 @@ +{ + "schema_version": 1, + "status": "local remediation acceptance gates passed", + "method": "bounded sequential local benchmark; no concurrency, network, or production access", + "fixture": { + "kind": "deterministic synthetic production-count fixture", + "materials": 1727, + "material_elements": 4144, + "material_applications": 128 + }, + "default_request": { + "material_id": 1, + "depth": 2, + "limit": 25, + "batch_calls": 2, + "expanded_materials": 25, + "database_query_count": 12, + "warm_median_service_ms": 342.218, + "payload_bytes": 9657, + "payload_sha256": "50f22e6a428bd19f7569f169361e71cf5d019056c23e07d75cf2af97a2f2fd11" + }, + "baseline_comparison": { + "database_query_count_before": 150, + "database_query_count_after": 12, + "query_reduction_percent": 92.0, + "warm_median_service_ms_before": 1638.338, + "warm_median_service_ms_after": 342.218, + "warm_speedup_multiple": 4.79, + "simulated_100_ms_rtt_end_to_end_ms_before": 16749.582, + "simulated_100_ms_rtt_end_to_end_ms_after": 1756.175, + "simulated_100_ms_rtt_speedup_multiple": 9.54 + }, + "limit_scaling": [ + {"limit": 1, "queries": 6, "warm_median_service_ms": 72.422, "simulated_100_ms_rtt_ms": 640.311, "payload_bytes": 336}, + {"limit": 5, "queries": 12, "warm_median_service_ms": 141.766, "simulated_100_ms_rtt_ms": 1447.459, "payload_bytes": 2043}, + {"limit": 10, "queries": 12, "warm_median_service_ms": 170.726, "simulated_100_ms_rtt_ms": 1456.539, "payload_bytes": 3945}, + {"limit": 25, "queries": 12, "warm_median_service_ms": 342.218, "simulated_100_ms_rtt_ms": 1756.175, "payload_bytes": 9657}, + {"limit": 50, "queries": 12, "warm_median_service_ms": 720.542, "simulated_100_ms_rtt_ms": 2077.637, "payload_bytes": 19262}, + {"limit": 100, "queries": 12, "warm_median_service_ms": 1720.358, "simulated_100_ms_rtt_ms": 2823.65, "payload_bytes": 37596} + ], + "semantic_stability": { + "default_payload_hash_matches_diagnostic_baseline": true, + "default_payload_bytes_match_diagnostic_baseline": true, + "repeated_results_identical": true, + "individual_and_batch_neighbor_results_equal": true + }, + "acceptance": { + "depth_two_queries_at_most_12": true, + "default_warm_median_at_most_500_ms": true, + "default_100_ms_rtt_at_most_3000_ms": true, + "serialization_below_5_ms": true, + "limits_50_and_100_below_20_seconds": true, + "disposable_postgresql_validation": "pending where available", + "production_validation": "not authorized and not performed" + } +} diff --git a/scripts/diagnose_neighborhood_performance.py b/scripts/diagnose_neighborhood_performance.py index 3ef483c..dbe8c3d 100644 --- a/scripts/diagnose_neighborhood_performance.py +++ b/scripts/diagnose_neighborhood_performance.py @@ -33,7 +33,6 @@ from app.services.material.neighbor_service import MaterialNeighborService from app.services.material.neighborhood_service import MaterialNeighborhoodService - ELEMENTS = ("Li", "Fe", "P", "O", "Na", "Mn", "Co", "Ni", "Si", "S") APPLICATIONS = ("battery", "catalyst", "structural") @@ -56,15 +55,17 @@ class ProfiledNeighborService(MaterialNeighborService): def __init__(self, db): super().__init__(db) self.calls = 0 + self.materials_requested = 0 self.call_ms = 0.0 self.build_ms = 0.0 - def get_neighbors(self, material_id: int) -> dict: + def get_neighbors_batch(self, material_ids: list[int]) -> dict[int, dict]: started = time.perf_counter() try: - return super().get_neighbors(material_id) + return super().get_neighbors_batch(material_ids) finally: self.calls += 1 + self.materials_requested += len(material_ids) self.call_ms += (time.perf_counter() - started) * 1_000 def _build_neighbors(self, *, neighbor_scores, materials_by_id): @@ -271,6 +272,7 @@ def service_measurement(depth: int, limit: int) -> tuple[dict[str, Any], dict]: "database_execute_ms": round(database_ms, 3), "query_count_by_table": dict(sorted(statement_counts.items())), "neighbor_service_calls": profiled.calls, + "neighbor_materials_requested": profiled.materials_requested, "neighbor_service_ms": round(profiled.call_ms, 3), "neighbor_build_and_score_ms": round(profiled.build_ms, 3), "neighborhood_graph_ms_estimate": round( @@ -300,11 +302,10 @@ def override_db(): app.dependency_overrides[get_db] = override_db path = f"/api/v1/materials/1/neighborhood?depth={depth}&limit={limit}" try: - with simulated_round_trip(engine, rtt_ms), record_queries(engine) as queries: - with TestClient(app) as client: - started = time.perf_counter() - response = client.get(path) - wall_ms = (time.perf_counter() - started) * 1_000 + with simulated_round_trip(engine, rtt_ms), record_queries(engine) as queries, TestClient(app) as client: + started = time.perf_counter() + response = client.get(path) + wall_ms = (time.perf_counter() - started) * 1_000 finally: app.dependency_overrides.clear() if response.status_code != 200: diff --git a/tests/services/material/test_neighbor_batch_service.py b/tests/services/material/test_neighbor_batch_service.py new file mode 100644 index 0000000..d111007 --- /dev/null +++ b/tests/services/material/test_neighbor_batch_service.py @@ -0,0 +1,143 @@ +from collections.abc import Iterator + +import pytest +from sqlalchemy import create_engine, event +from sqlalchemy.orm import Session + +from app.core.database import Base +from app.models.application import Application +from app.models.element import Element +from app.models.material import Material +from app.models.material_application import MaterialApplication +from app.models.material_element import MaterialElement +from app.services.material.neighbor_service import MaterialNeighborService + + +@pytest.fixture +def neighbor_db() -> Iterator[Session]: + engine = create_engine("sqlite:///:memory:") + tables = [ + Material.__table__, + Element.__table__, + Application.__table__, + MaterialElement.__table__, + MaterialApplication.__table__, + ] + Base.metadata.create_all(engine, tables=tables) + + with Session(engine) as db: + db.add_all( + [ + Material( + id=1, + mp_id="mp-1", + formula="AB", + pretty_formula="AB", + material_type="test", + is_stable=True, + energy_above_hull=0.0, + ), + Material( + id=2, + mp_id="mp-2", + formula="AC", + pretty_formula="AC", + material_type=None, + is_stable=False, + energy_above_hull=None, + ), + Material( + id=3, + mp_id="mp-3", + formula="BC", + pretty_formula="BC", + material_type="test", + is_stable=True, + energy_above_hull=0.1, + ), + Material( + id=4, + mp_id="mp-4", + formula="D", + pretty_formula="D", + material_type="isolated", + is_stable=False, + energy_above_hull=None, + ), + Element(id=1, symbol="A", name="A"), + Element(id=2, symbol="B", name="B"), + Element(id=3, symbol="C", name="C"), + Element(id=4, symbol="D", name="D"), + Application(id=1, name="storage"), + Application(id=2, name="catalysis"), + ] + ) + db.flush() + db.add_all( + [ + MaterialElement(material_id=1, element_id=1, fraction=0.5), + MaterialElement(material_id=1, element_id=2, fraction=0.5), + MaterialElement(material_id=2, element_id=1, fraction=0.5), + MaterialElement(material_id=2, element_id=3, fraction=0.5), + MaterialElement(material_id=3, element_id=2, fraction=0.5), + MaterialElement(material_id=3, element_id=3, fraction=0.5), + MaterialElement(material_id=4, element_id=4, fraction=1.0), + MaterialApplication(material_id=1, application_id=1), + MaterialApplication(material_id=2, application_id=1), + MaterialApplication(material_id=2, application_id=2), + MaterialApplication(material_id=3, application_id=2), + ] + ) + db.commit() + yield db + + engine.dispose() + + +def test_batch_neighbors_are_exactly_equal_to_individual_results( + neighbor_db: Session, +) -> None: + service = MaterialNeighborService(neighbor_db) + material_ids = [1, 2, 3, 4, 999] + expected = { + material_id: service.get_neighbors(material_id) + for material_id in material_ids + } + + assert service.get_neighbors_batch(material_ids) == expected + + +def test_batch_neighbor_query_count_is_bounded(neighbor_db: Session) -> None: + service = MaterialNeighborService(neighbor_db) + query_count = 0 + + def count_query(*_args, **_kwargs) -> None: + nonlocal query_count + query_count += 1 + + event.listen( + neighbor_db.get_bind(), + "before_cursor_execute", + count_query, + ) + try: + service.get_neighbors_batch([1, 2, 3, 4]) + finally: + event.remove( + neighbor_db.get_bind(), + "before_cursor_execute", + count_query, + ) + + assert query_count <= 6 + + +def test_batch_neighbors_deduplicate_requested_material_ids( + neighbor_db: Session, +) -> None: + service = MaterialNeighborService(neighbor_db) + + assert service.get_neighbors_batch([2, 1, 2]) == { + 2: service.get_neighbors(2), + 1: service.get_neighbors(1), + } diff --git a/tests/services/material/test_neighborhood_service.py b/tests/services/material/test_neighborhood_service.py index 8ebd693..276b393 100644 --- a/tests/services/material/test_neighborhood_service.py +++ b/tests/services/material/test_neighborhood_service.py @@ -68,8 +68,11 @@ def _service() -> MaterialNeighborhoodService: 4: _response(4, []), } - service.neighbor_service.get_neighbors = Mock( - side_effect=lambda material_id: responses[material_id] + service.neighbor_service.get_neighbors_batch = Mock( + side_effect=lambda material_ids: { + material_id: responses[material_id] + for material_id in material_ids + } ) return service @@ -154,9 +157,9 @@ def test_limit_bounds_neighbor_expansion() -> None: for node in result["nodes"] } == {1, 2} - assert service.neighbor_service.get_neighbors.call_args_list == [ - call(1), - call(2), + assert service.neighbor_service.get_neighbors_batch.call_args_list == [ + call([1]), + call([2]), ] @@ -170,8 +173,8 @@ def test_limit_one_does_not_expand_descendants() -> None: ) assert [node["material_id"] for node in result["nodes"]] == [1] - assert service.neighbor_service.get_neighbors.call_count == 1 - service.neighbor_service.get_neighbors.assert_called_once_with(1) + assert service.neighbor_service.get_neighbors_batch.call_count == 1 + service.neighbor_service.get_neighbors_batch.assert_called_once_with([1]) def test_bounded_traversal_is_deterministic() -> None: @@ -204,8 +207,11 @@ def build_service(root_neighbor_ids: list[int]) -> MaterialNeighborhoodService: 2: _response(2, []), 3: _response(3, []), } - service.neighbor_service.get_neighbors = Mock( - side_effect=lambda material_id: responses[material_id] + service.neighbor_service.get_neighbors_batch = Mock( + side_effect=lambda material_ids: { + material_id: responses[material_id] + for material_id in material_ids + } ) return service @@ -221,15 +227,13 @@ def build_service(root_neighbor_ids: list[int]) -> MaterialNeighborhoodService: (edge["source_material_id"], edge["target_material_id"]) for edge in first["edges"] ] == [(1, 2), (1, 3)] - assert first_service.neighbor_service.get_neighbors.call_args_list == [ - call(1), - call(2), - call(3), + assert first_service.neighbor_service.get_neighbors_batch.call_args_list == [ + call([1]), + call([2, 3]), ] - assert second_service.neighbor_service.get_neighbors.call_args_list == [ - call(1), - call(2), - call(3), + assert second_service.neighbor_service.get_neighbors_batch.call_args_list == [ + call([1]), + call([2, 3]), ] @@ -247,5 +251,5 @@ def test_depth_remains_maximum_expansion_depth() -> None: for node in result["nodes"] } == {1, 2, 3} - assert service.neighbor_service.get_neighbors.call_count == 1 - service.neighbor_service.get_neighbors.assert_called_once_with(1) \ No newline at end of file + assert service.neighbor_service.get_neighbors_batch.call_count == 1 + service.neighbor_service.get_neighbors_batch.assert_called_once_with([1])