deepnote-toolkit/tests/unit/test_set_integrations_env.py at ed1621b03ed8aff0247474572545590840ef9b39 · deepnote/deepnote-toolkit · 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
import types
from unittest import mock

import pytest

from deepnote_toolkit import env as dnenv
from deepnote_toolkit.set_integrations_env import set_integration_env


def _mock_cfg(enabled: bool) -> types.SimpleNamespace:
    return types.SimpleNamespace(
        runtime=types.SimpleNamespace(env_integration_enabled=enabled)
    )


@mock.patch(
    "deepnote_toolkit.set_integrations_env.get_config",
    side_effect=lambda: _mock_cfg(True),
)
def test_set_integration_env_success(mock_get_config, monkeypatch):  # noqa: ARG001
    # Fake session with two variables
    class DummyResp:
        ok = True

        def json(self):
            return [{"name": "X_A", "value": "1"}, {"name": "X_B", "value": "2"}]

    class DummySession:
        def mount(self, *a, **k):
            pass

        def get(self, *a, **k):
            return DummyResp()

    monkeypatch.setattr(
        "deepnote_toolkit.set_integrations_env.requests.Session", lambda: DummySession()
    )
    monkeypatch.setenv("DEEPNOTE_PROJECT_ID", "pid")
    # Clean env
    dnenv.unset_env("X_A")
    dnenv.unset_env("X_B")
    try:
        set_integration_env()
        assert dnenv.get_env("X_A") == "1"
        assert dnenv.get_env("X_B") == "2"
    finally:
        dnenv.unset_env("X_A")
        dnenv.unset_env("X_B")


@mock.patch(
    "deepnote_toolkit.set_integrations_env.get_config",
    side_effect=lambda: _mock_cfg(True),
)
def test_set_integration_env_http_error(mock_get_config, monkeypatch):  # noqa: ARG001
    class DummyResp:
        ok = False

        def json(self):
            return []

    class DummySession:
        def mount(self, *a, **k):
            pass

        def get(self, *a, **k):
            return DummyResp()

    monkeypatch.setattr(
        "deepnote_toolkit.set_integrations_env.requests.Session", lambda: DummySession()
    )
    monkeypatch.setenv("DEEPNOTE_PROJECT_ID", "pid")
    with pytest.raises(Exception, match="Failed to fetch integration variables"):
        set_integration_env()


@mock.patch(
    "deepnote_toolkit.set_integrations_env.get_config",
    side_effect=lambda: _mock_cfg(True),
)
def test_set_integration_env_clears_config_cache(
    mock_get_config, monkeypatch  # noqa: ARG001
):
    class DummyResp:
        ok = True

        def json(self):
            return [{"name": "DEEPNOTE_DO_NOT_COERCE_FLOAT", "value": "1"}]

    class DummySession:
        def mount(self, *a, **k):
            pass

        def get(self, *a, **k):
            return DummyResp()

    monkeypatch.setattr(
        "deepnote_toolkit.set_integrations_env.requests.Session", lambda: DummySession()
    )
    monkeypatch.setenv("DEEPNOTE_PROJECT_ID", "pid")

    with mock.patch(
        "deepnote_toolkit.set_integrations_env.clear_config_cache"
    ) as mock_clear:
        set_integration_env()
        mock_clear.assert_called_once()


def test_set_integration_env_disabled_gate(monkeypatch):
    # get_config returns disabled gate; Session should not be constructed
    from deepnote_toolkit import set_integrations_env as sie

    monkeypatch.setattr(sie, "get_config", lambda: _mock_cfg(False))

    class Boom:
        def __init__(self, *a, **k):  # noqa: ARG002
            raise AssertionError("Session should not be constructed when disabled")

    monkeypatch.setattr(sie.requests, "Session", Boom)
    # Ensure no env is set and no exception raised
    prev = dnenv.get_env("SHOULD_NOT_EXIST")
    sie.set_integration_env()
    assert dnenv.get_env("SHOULD_NOT_EXIST") == prev