infrahub-sdk-python/tests/unit/sdk/test_store.py at 42ac2d4fcd96e1d78be4604a72002c3419edbfc0 · opsmill/infrahub-sdk-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
from __future__ import annotations

from typing import TYPE_CHECKING

import pytest

from infrahub_sdk.exceptions import NodeInvalidError, NodeNotFoundError
from infrahub_sdk.node import InfrahubNode, InfrahubNodeSync
from infrahub_sdk.protocols import BuiltinIPAddressSync, BuiltinIPPrefix
from infrahub_sdk.store import NodeStore, NodeStoreSync, get_schema_name

if TYPE_CHECKING:
    from infrahub_sdk.schema import NodeSchemaAPI

    from .conftest import BothClients

client_types = ["standard", "sync"]


@pytest.mark.parametrize("client_type", client_types)
def test_node_store_set(client_type: str, clients: BothClients, schema_with_hfid: dict[str, NodeSchemaAPI]) -> None:
    if client_type == "standard":
        client = clients.standard
        store = NodeStore(default_branch="main")
        node_class = InfrahubNode
    else:
        client = clients.sync
        store = NodeStoreSync(default_branch="main")
        node_class = InfrahubNodeSync

    data = {
        "name": {"value": "JFK1"},
        "description": {"value": "JFK Airport"},
        "type": {"value": "SITE"},
    }
    node = node_class(client=client, schema=schema_with_hfid["location"], data=data)

    store.set(key="mykey", node=node)

    assert node._internal_id in store._branches[client.default_branch]._objs
    assert "mykey" in store._branches[client.default_branch]._keys
    assert store._branches[client.default_branch]._keys["mykey"] == node._internal_id


@pytest.mark.parametrize("client_type", client_types)
def test_node_store_set_no_hfid(client_type: str, clients: BothClients, location_schema: NodeSchemaAPI) -> None:
    if client_type == "standard":
        client = clients.standard
        store = NodeStore(default_branch="main")
        node_class = InfrahubNode
    else:
        client = clients.sync
        store = NodeStoreSync(default_branch="main")
        node_class = InfrahubNodeSync

    data = {
        "name": {"value": "JFK1"},
        "description": {"value": "JFK Airport"},
        "type": {"value": "SITE"},
    }
    node = node_class(client=client, schema=location_schema, data=data)

    store.set(key="mykey", node=node)

    assert store._branches[client.default_branch]._objs[node._internal_id] == node
    assert node._internal_id in store._branches[client.default_branch]._objs
    assert "mykey" in store._branches[client.default_branch]._keys
    assert store._branches[client.default_branch]._keys["mykey"] == node._internal_id

    branch_name = "mybranch"
    store.set(key="anotherkey", node=node, branch=branch_name)

    assert store._branches[branch_name]._objs[node._internal_id] == node
    assert node._internal_id in store._branches[branch_name]._objs
    assert "anotherkey" in store._branches[branch_name]._keys
    assert store._branches[branch_name]._keys["anotherkey"] == node._internal_id


@pytest.mark.parametrize("client_type", client_types)
def test_node_store_get(client_type: str, clients: BothClients, location_schema: NodeSchemaAPI) -> None:
    if client_type == "standard":
        client = clients.standard
        store = NodeStore(default_branch="main")
        node_class = InfrahubNode
    else:
        client = clients.sync
        store = NodeStoreSync(default_branch="main")
        node_class = InfrahubNodeSync

    data = {
        "id": "54f3108c-1f21-44c4-93cf-ec5737587b48",
        "name": {"value": "JFK1"},
        "description": {"value": "JFK Airport"},
        "type": {"value": "SITE"},
    }
    node = node_class(client=client, schema=location_schema, data=data)

    store.set(key="mykey", node=node)

    assert store.get(key=node._internal_id).id == node.id
    assert store.get(kind="BuiltinLocation", key="mykey").id == node.id
    assert store.get(key="mykey").id == node.id

    assert store.get(kind="BuiltinTest", key="mykey", raise_when_missing=False) is None

    assert store.get(kind="BuiltinLocation", key="anotherkey", raise_when_missing=False) is None
    assert store.get(key="anotherkey", raise_when_missing=False) is None

    with pytest.raises(NodeNotFoundError):
        store.get(kind="BuiltinLocation", key="anotherkey")
    with pytest.raises(NodeNotFoundError):
        store.get(key="anotherkey")
    with pytest.raises(NodeNotFoundError):
        store.get(key="mykey", branch="mybranch")

    with pytest.raises(NodeInvalidError):
        store.get(kind="BuiltinTest", key="mykey")

    store.set(key="mykey", node=node, branch="mybranch")
    assert store.get(key="mykey", branch="mybranch").id == node.id


@pytest.mark.parametrize("client_type", client_types)
def test_node_store_get_with_hfid(
    client_type: str, clients: BothClients, schema_with_hfid: dict[str, NodeSchemaAPI]
) -> None:
    if client_type == "standard":
        client = clients.standard
        store = NodeStore(default_branch="main")
        node_class = InfrahubNode
    else:
        client = clients.sync
        store = NodeStoreSync(default_branch="main")
        node_class = InfrahubNodeSync

    data = {
        "id": "54f3108c-1f21-44c4-93cf-ec5737587b48",
        "name": {"value": "JFK1"},
        "description": {"value": "JFK Airport"},
        "type": {"value": "SITE"},
    }
    node = node_class(client=client, schema=schema_with_hfid["location"], data=data)

    store.set(key="mykey", node=node)

    assert store.get(key=node._internal_id).id == node.id

    assert store.get(kind="BuiltinLocation", key="mykey").id == node.id
    assert store.get(key="mykey").id == node.id
    assert store.get(key="BuiltinLocation__JFK1").id == node.id
    assert store.get(kind="BuiltinLocation", key="JFK1").id == node.id
    assert store.get(key="54f3108c-1f21-44c4-93cf-ec5737587b48").id == node.id

    assert store.get(kind="BuiltinLocation", key="anotherkey", raise_when_missing=False) is None
    assert store.get(key="anotherkey", raise_when_missing=False) is None

    with pytest.raises(NodeNotFoundError):
        store.get(kind="BuiltinLocation", key="anotherkey")
    with pytest.raises(NodeNotFoundError):
        store.get(key="anotherkey")


def test_store_get_schema_name() -> None:
    assert get_schema_name(schema=BuiltinIPPrefix) == BuiltinIPPrefix.__name__
    assert get_schema_name(schema=BuiltinIPAddressSync) == BuiltinIPAddressSync.__name__[:-4]