wildedge-python/tests/test_offline_replay.py at 0db3079e141bb316431a6ae2f9980db0fc417836 · wild-edge/wildedge-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
from __future__ import annotations

from unittest.mock import patch

from wildedge.client import WildEdge
from wildedge.device import DeviceInfo


class _DummyConsumer:
    def __init__(self, *args, **kwargs):
        pass

    def flush(self, timeout: float = 5.0) -> None:
        pass

    def close(self, timeout: float | None = None) -> None:
        pass

    def _pause(self) -> None:
        pass

    def _resume(self) -> None:
        pass


class _Model:
    pass


def test_offline_replay_restores_model_registry_for_pending_events(tmp_path):
    queue_dir = tmp_path / "queue"
    dead_dir = tmp_path / "dead"
    with (
        patch(
            "wildedge.client.detect_device",
            return_value=DeviceInfo(device_id="d", device_type="linux"),
        ),
        patch("wildedge.client.Transmitter"),
        patch("wildedge.client.Consumer", _DummyConsumer),
    ):
        client_a = WildEdge(
            dsn="https://secret@ingest.wildedge.dev/proj",
            app_identity="app-a",
            offline_queue_dir=str(queue_dir),
            dead_letter_dir=str(dead_dir),
            enable_offline_persistence=True,
        )
        client_a.register_model(
            _Model(),
            model_id="ResNet",
            source="local",
            family="resnet",
            version="1.0",
            quantization="fp32",
        )
        client_a.publish(
            {"event_id": "e1", "event_type": "model_load", "model_id": "ResNet"}
        )
        client_a.close()

        client_b = WildEdge(
            dsn="https://secret@ingest.wildedge.dev/proj",
            app_identity="app-a",
            offline_queue_dir=str(queue_dir),
            dead_letter_dir=str(dead_dir),
            enable_offline_persistence=True,
        )

    assert client_b.queue.length() == 1
    models = client_b.registry.snapshot()
    assert "ResNet" in models
    assert models["ResNet"]["model_name"] == "_Model"