refactor ipam unit tests to maintain data (#5163)

* add module-level backend fixtures for IPAM unit tests

* support memgraph

* move fixtures around a little more

* diff cleanup and os.path cleanup

* refactor to reduce duplication

backend/infrahub/core/query/delete.py

@@ -0,0 +1,58 @@
+from typing import Any
+
+from infrahub import config
+from infrahub.core.query import Query, QueryType
+from infrahub.core.timestamp import Timestamp
+from infrahub.database import InfrahubDatabase
+
+
+class DeleteAfterTimeQuery(Query):
+    name: str = "delete_after_time"
+    insert_return: bool = False
+    type: QueryType = QueryType.WRITE
+
+    def __init__(self, timestamp: Timestamp, **kwargs: Any):
+        self.timestamp = timestamp
+        super().__init__(**kwargs)
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:
+        self.params = {"timestamp": self.timestamp.to_string()}
+        query_1 = """
+        // ---------------------
+        // Reset edges with to time after timestamp
+        // ---------------------
+        CALL {
+            OPTIONAL MATCH (p)-[r]-(q)
+            WHERE r.to > $timestamp
+            SET r.to = NULL
+        }
+        """
+        self.add_to_query(query_1)
+        if config.SETTINGS.database.db_type is config.DatabaseType.NEO4J:
+            query_2 = """
+            // ---------------------
+            // Delete edges with from time after timestamp timestamp
+            // ---------------------
+            CALL {
+                OPTIONAL MATCH (p)-[r]->(q)
+                WHERE r.from > $timestamp
+                DELETE r
+                WITH p, q
+                UNWIND [p, q] AS maybe_orphan
+                WITH maybe_orphan
+                WHERE NOT exists((maybe_orphan)--())
+                DELETE maybe_orphan
+            }
+            """
+        else:
+            query_2 = """
+            // ---------------------
+            // Delete edges with from time after timestamp timestamp
+            // ---------------------
+            CALL {
+                OPTIONAL MATCH (p)-[r]->(q)
+                WHERE r.from > $timestamp
+                DELETE r
+            }
+            """
+        self.add_to_query(query_2)

backend/tests/conftest.py

@@ -125,17 +125,29 @@ async def db(
 
 @pytest.fixture
 async def empty_database(db: InfrahubDatabase) -> None:
+    await do_empty_database(db=db)
+
+
+async def do_empty_database(db: InfrahubDatabase) -> None:
     await delete_all_nodes(db=db)
     await create_root_node(db=db)
 
 
 @pytest.fixture
 async def reset_registry(db: InfrahubDatabase) -> None:
+    await do_reset_registry(db=db)
+
+
+async def do_reset_registry(db: InfrahubDatabase) -> None:
     registry.delete_all()
 
 
 @pytest.fixture
 async def default_branch(reset_registry, local_storage_dir, empty_database, db: InfrahubDatabase) -> Branch:
+    return await do_default_branch(db=db)
+
+
+async def do_default_branch(db: InfrahubDatabase) -> Branch:
     branch = await create_default_branch(db=db)
     await create_global_branch(db=db)
     registry.schema = SchemaManager()
@@ -153,6 +165,10 @@ async def default_ipnamespace(db: InfrahubDatabase, register_core_models_schema)
 
 @pytest.fixture
 def local_storage_dir(tmp_path: Path) -> Path:
+    return do_local_storage_dir(tmp_path=tmp_path)
+
+
+def do_local_storage_dir(tmp_path: Path) -> Path:
     storage_dir = tmp_path / "storage"
     storage_dir.mkdir()
 
@@ -164,17 +180,25 @@ def local_storage_dir(tmp_path: Path) -> Path:
 
 @pytest.fixture
 async def register_internal_models_schema(default_branch: Branch) -> SchemaBranch:
+    return await do_register_internal_models_schema(branch=default_branch)
+
+
+async def do_register_internal_models_schema(branch: Branch) -> SchemaBranch:
     schema = SchemaRoot(**internal_schema)
-    schema_branch = registry.schema.register_schema(schema=schema, branch=default_branch.name)
-    default_branch.update_schema_hash()
+    schema_branch = registry.schema.register_schema(schema=schema, branch=branch.name)
+    branch.update_schema_hash()
     return schema_branch
 
 
 @pytest.fixture
 async def register_core_models_schema(default_branch: Branch, register_internal_models_schema) -> SchemaBranch:
+    return await do_register_core_models_schema(branch=default_branch)
+
+
+async def do_register_core_models_schema(branch: Branch) -> SchemaBranch:
     schema = SchemaRoot(**core_models)
-    schema_branch = registry.schema.register_schema(schema=schema, branch=default_branch.name)
-    default_branch.update_schema_hash()
+    schema_branch = registry.schema.register_schema(schema=schema, branch=branch.name)
+    branch.update_schema_hash()
     return schema_branch
 
 

backend/tests/unit/conftest.py

@@ -2413,7 +2413,7 @@ async def builtin_schema() -> SchemaRoot:
     return SchemaRoot(**SCHEMA)
 
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 async def ipam_schema() -> SchemaRoot:
     SCHEMA: dict[str, Any] = {
         "nodes": [

backend/tests/unit/core/ipam/conftest.py

@@ -1,17 +1,82 @@
+from pathlib import Path
+
 import pytest
 
 from infrahub.core import registry
 from infrahub.core.branch import Branch
 from infrahub.core.constants import InfrahubKind
+from infrahub.core.diff.repository.repository import DiffRepository
+from infrahub.core.initialization import (
+    create_ipam_namespace,
+)
 from infrahub.core.node import Node
+from infrahub.core.query.delete import DeleteAfterTimeQuery
+from infrahub.core.schema import SchemaRoot
 from infrahub.core.schema.schema_branch import SchemaBranch
+from infrahub.core.timestamp import Timestamp
 from infrahub.database import InfrahubDatabase
+from infrahub.dependencies.registry import get_component_registry
+from tests.conftest import (
+    do_default_branch,
+    do_empty_database,
+    do_local_storage_dir,
+    do_register_core_models_schema,
+    do_register_internal_models_schema,
+    do_reset_registry,
+)
 
 
-@pytest.fixture
-async def ip_dataset_01(
+@pytest.fixture(scope="module")
+async def register_internal_models_schema(default_branch: Branch) -> SchemaBranch:
+    return await do_register_internal_models_schema(branch=default_branch)
+
+
+@pytest.fixture(scope="module")
+async def register_core_models_schema(default_branch: Branch, register_internal_models_schema) -> SchemaBranch:
+    return await do_register_core_models_schema(branch=default_branch)
+
+
+@pytest.fixture(scope="module")
+def local_storage_dir(tmp_path_module_scope: Path) -> Path:
+    return do_local_storage_dir(tmp_path=tmp_path_module_scope)
+
+
+@pytest.fixture(scope="module")
+async def empty_database(db: InfrahubDatabase) -> None:
+    await do_empty_database(db=db)
+
+
+@pytest.fixture(scope="module")
+async def reset_registry(db: InfrahubDatabase) -> None:
+    await do_reset_registry(db=db)
+
+
+@pytest.fixture(scope="module")
+async def default_branch(reset_registry, local_storage_dir, empty_database, db: InfrahubDatabase) -> Branch:
+    return await do_default_branch(db=db)
+
+
+@pytest.fixture(scope="module")
+async def default_ipnamespace(db: InfrahubDatabase, register_core_models_schema) -> Node | None:
+    if not registry._default_ipnamespace:
+        ip_namespace = await create_ipam_namespace(db=db)
+        registry.default_ipnamespace = ip_namespace.id
+        return ip_namespace
+    return None
+
+
+@pytest.fixture(scope="module")
+async def register_ipam_schema(default_branch: Branch, ipam_schema: SchemaRoot) -> SchemaBranch:
+    schema_branch = registry.schema.register_schema(schema=ipam_schema, branch=default_branch.name)
+    default_branch.update_schema_hash()
+    return schema_branch
+
+
+@pytest.fixture(scope="module")
+async def ip_dataset_01_load(
     db: InfrahubDatabase,
     default_branch: Branch,
+    default_ipnamespace: Node,
     register_core_models_schema: SchemaBranch,
     register_ipam_schema: SchemaBranch,
 ):
@@ -103,3 +168,35 @@ async def ip_dataset_01(
         "net242": net242,
     }
     return data
+
+
+@pytest.fixture(scope="module")
+async def diff_repository(db: InfrahubDatabase, default_branch: Branch) -> DiffRepository:
+    component_registry = get_component_registry()
+    return await component_registry.get_component(DiffRepository, db=db, branch=default_branch)
+
+
+@pytest.fixture(scope="module")
+def start_time():
+    return Timestamp()
+
+
+@pytest.fixture(scope="function")
+async def ip_dataset_01(
+    db: InfrahubDatabase,
+    default_branch: Branch,
+    ip_dataset_01_load,
+    diff_repository: DiffRepository,
+    start_time: Timestamp,
+):
+    yield ip_dataset_01_load
+
+    all_diff_roots = await diff_repository.get_empty_roots()
+    root_uuids_to_delete = []
+    for diff_root in all_diff_roots:
+        if start_time <= diff_root.from_time:
+            root_uuids_to_delete.append(diff_root.uuid)
+    await diff_repository.delete_diff_roots(diff_root_uuids=root_uuids_to_delete)
+
+    query = await DeleteAfterTimeQuery.init(db=db, timestamp=start_time)
+    await query.execute(db=db)

backend/tests/unit/core/ipam/test_ipam_reconcile_query.py

@@ -3,7 +3,7 @@
 from infrahub.core import registry
 from infrahub.core.branch import Branch
 from infrahub.core.constants import InfrahubKind
-from infrahub.core.initialization import create_branch, create_ipam_namespace, get_default_ipnamespace
+from infrahub.core.initialization import create_branch, get_default_ipnamespace
 from infrahub.core.manager import NodeManager
 from infrahub.core.node import Node
 from infrahub.core.query.ipam import IPPrefixReconcileQuery
@@ -12,7 +12,6 @@
 
 
 async def test_ipprefix_reconcile_query_simple(db: InfrahubDatabase, default_branch: Branch, ip_dataset_01):
-    await create_ipam_namespace(db=db)
     default_ipnamespace = await get_default_ipnamespace(db=db)
     registry.default_ipnamespace = default_ipnamespace.id
     prefix_140 = ip_dataset_01["net140"]

backend/tests/unit/core/ipam/test_ipam_reconciler.py

@@ -4,28 +4,23 @@
 
 from infrahub.core import registry
 from infrahub.core.branch import Branch
-from infrahub.core.initialization import create_ipam_namespace, get_default_ipnamespace
+from infrahub.core.initialization import (
+    get_default_ipnamespace,
+)
 from infrahub.core.ipam.reconciler import IpamReconciler
 from infrahub.core.manager import NodeManager
 from infrahub.core.node import Node
 from infrahub.database import InfrahubDatabase
 from infrahub.exceptions import NodeNotFoundError
 
 
-async def test_invalid_ip_node_raises_error(db: InfrahubDatabase, default_branch: Branch, register_core_models_schema):
-    await create_ipam_namespace(db=db)
-    default_ipnamespace = await get_default_ipnamespace(db=db)
-
-    reconciler = IpamReconciler(db=db, branch=default_branch)
-    with pytest.raises(NodeNotFoundError):
-        await reconciler.reconcile(ip_value=ipaddress.ip_interface("192.168.1.1"), namespace=default_ipnamespace)
-
-
 async def test_first_prefix(
-    db: InfrahubDatabase, default_branch: Branch, register_core_models_schema, register_ipam_schema
+    db: InfrahubDatabase,
+    default_branch: Branch,
+    register_core_models_schema,
+    register_ipam_schema,
+    default_ipnamespace: Node,
 ):
-    await create_ipam_namespace(db=db)
-    default_ipnamespace = await get_default_ipnamespace(db=db)
     prefix_schema = registry.schema.get_node_schema(name="IpamIPPrefix", branch=default_branch)
     net161 = await Node.init(db=db, schema=prefix_schema)
     await net161.new(db=db, prefix="2001:db8::/48", ip_namespace=default_ipnamespace)
@@ -40,8 +35,15 @@ async def test_first_prefix(
     assert all_prefixes[0].is_top_level.value is True
 
 
+async def test_invalid_ip_node_raises_error(db: InfrahubDatabase, default_branch: Branch, ip_dataset_01):
+    default_ipnamespace = await get_default_ipnamespace(db=db)
+
+    reconciler = IpamReconciler(db=db, branch=default_branch)
+    with pytest.raises(NodeNotFoundError):
+        await reconciler.reconcile(ip_value=ipaddress.ip_interface("192.168.1.1"), namespace=default_ipnamespace)
+
+
 async def test_ipprefix_reconciler_no_change(db: InfrahubDatabase, default_branch: Branch, ip_dataset_01):
-    await create_ipam_namespace(db=db)
     default_ipnamespace = await get_default_ipnamespace(db=db)
     registry.default_ipnamespace = default_ipnamespace.id
     prefix_140 = ip_dataset_01["net140"]
@@ -63,7 +65,6 @@ async def test_ipprefix_reconciler_no_change(db: InfrahubDatabase, default_branc
 
 
 async def test_ipprefix_reconciler_new_prefix_update(db: InfrahubDatabase, default_branch: Branch, ip_dataset_01):
-    await create_ipam_namespace(db=db)
     default_ipnamespace = await get_default_ipnamespace(db=db)
     registry.default_ipnamespace = default_ipnamespace.id
     prefix_schema = registry.schema.get_node_schema(name="IpamIPPrefix", branch=default_branch)
@@ -114,7 +115,6 @@ async def test_ipprefix_reconciler_new_prefix_update(db: InfrahubDatabase, defau
 
 
 async def test_ipprefix_reconciler_new_address_update(db: InfrahubDatabase, default_branch: Branch, ip_dataset_01):
-    await create_ipam_namespace(db=db)
     default_ipnamespace = await get_default_ipnamespace(db=db)
     registry.default_ipnamespace = default_ipnamespace.id
     address_schema = registry.schema.get_node_schema(name="IpamIPAddress", branch=default_branch)
@@ -140,7 +140,6 @@ async def test_ipprefix_reconciler_new_address_update(db: InfrahubDatabase, defa
 
 
 async def test_ip_prefix_reconciler_delete_prefix(db: InfrahubDatabase, default_branch: Branch, ip_dataset_01):
-    await create_ipam_namespace(db=db)
     default_ipnamespace = await get_default_ipnamespace(db=db)
     registry.default_ipnamespace = default_ipnamespace.id
     namespace = ip_dataset_01["ns1"]
@@ -179,7 +178,6 @@ async def test_ip_prefix_reconciler_delete_prefix(db: InfrahubDatabase, default_
 
 
 async def test_ip_prefix_reconciler_delete_address(db: InfrahubDatabase, default_branch: Branch, ip_dataset_01):
-    await create_ipam_namespace(db=db)
     default_ipnamespace = await get_default_ipnamespace(db=db)
     registry.default_ipnamespace = default_ipnamespace.id
     namespace = ip_dataset_01["ns1"]
@@ -203,7 +201,6 @@ async def test_ip_prefix_reconciler_delete_address(db: InfrahubDatabase, default
 
 
 async def test_ipprefix_reconciler_prefix_value_update(db: InfrahubDatabase, default_branch: Branch, ip_dataset_01):
-    await create_ipam_namespace(db=db)
     default_ipnamespace = await get_default_ipnamespace(db=db)
     registry.default_ipnamespace = default_ipnamespace.id
     namespace = ip_dataset_01["ns1"]