cleveragents-core/benchmarks/bench_uko_persistence.py

"""ASV benchmarks for UKO Graph Persistence service.

Benchmarks the save() and restore() methods under varying triple counts.
"""

from __future__ import annotations

import tempfile
from pathlib import Path
from unittest.mock import MagicMock

from cleveragents.application.services.uko_persistence import (
    InMemoryPersistenceBackend,
    JSONFilePersistenceBackend,
    UKOGraphPersistence,
)
from cleveragents.domain.models.acms.index_backends import GraphIndexBackend


class UKOPersistenceBenchmarks:
    """Benchmarks for UKO Graph Persistence."""

    params = [10, 100, 1000]
    param_names = ["triple_count"]

    def setup(self, triple_count: int) -> None:
        """Set up benchmark fixtures."""
        self.triple_count = triple_count
        self.temp_dir = tempfile.mkdtemp()

        # Create mock graph backend
        self.graph_backend = MagicMock(spec=GraphIndexBackend)

        # Create sample triples
        self.triples = [
            {
                "subject": f"uko://code/module/mod_{i}",
                "predicate": "uko:type",
                "object": "uko-py:Module",
            }
            for i in range(triple_count)
        ]

        # Set up graph backend to return triples
        self.graph_backend.query.return_value = [
            {
                "s": t["subject"],
                "p": t["predicate"],
                "o": t["object"],
            }
            for t in self.triples
        ]

    def teardown(self, triple_count: int) -> None:
        """Clean up after benchmark."""
        import shutil

        shutil.rmtree(self.temp_dir, ignore_errors=True)

    def time_save_with_in_memory_backend(self, triple_count: int) -> None:
        """Benchmark save() with InMemoryPersistenceBackend."""
        persistence_backend = InMemoryPersistenceBackend()
        service = UKOGraphPersistence(
            graph_backend=self.graph_backend,
            project="local/benchmark",
            persistence_backend=persistence_backend,
        )
        service.save()

    def time_save_with_json_file_backend(self, triple_count: int) -> None:
        """Benchmark save() with JSONFilePersistenceBackend."""
        persistence_backend = JSONFilePersistenceBackend(self.temp_dir)
        service = UKOGraphPersistence(
            graph_backend=self.graph_backend,
            project="local/benchmark",
            persistence_backend=persistence_backend,
        )
        service.save()

    def time_restore_with_in_memory_backend(self, triple_count: int) -> None:
        """Benchmark restore() with InMemoryPersistenceBackend."""
        persistence_backend = InMemoryPersistenceBackend()
        persistence_backend.save("local/benchmark", self.triples)

        service = UKOGraphPersistence(
            graph_backend=self.graph_backend,
            project="local/benchmark",
            persistence_backend=persistence_backend,
        )
        service.restore()

    def time_restore_with_json_file_backend(self, triple_count: int) -> None:
        """Benchmark restore() with JSONFilePersistenceBackend."""
        persistence_backend = JSONFilePersistenceBackend(self.temp_dir)
        persistence_backend.save("local/benchmark", self.triples)

        service = UKOGraphPersistence(
            graph_backend=self.graph_backend,
            project="local/benchmark",
            persistence_backend=persistence_backend,
        )
        service.restore()

    def time_json_file_backend_load(self, triple_count: int) -> None:
        """Benchmark JSONFilePersistenceBackend.load()."""
        persistence_backend = JSONFilePersistenceBackend(self.temp_dir)
        persistence_backend.save("local/benchmark", self.triples)
        persistence_backend.load("local/benchmark")

    def time_in_memory_backend_load(self, triple_count: int) -> None:
        """Benchmark InMemoryPersistenceBackend.load()."""
        persistence_backend = InMemoryPersistenceBackend()
        persistence_backend.save("local/benchmark", self.triples)
        persistence_backend.load("local/benchmark")