prefect/integration-tests/test_worker.py at 0b734dbc89f10cc6e4e99887178fa01e8bb948eb · PrefectHQ/prefect · 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
import asyncio
import subprocess
import sys
from threading import Thread
from typing import List
from uuid import uuid4

import uv

from prefect.events import Event
from prefect.events.clients import get_events_subscriber
from prefect.events.filters import EventFilter, EventNameFilter, EventOccurredFilter
from prefect.types._datetime import now


async def watch_worker_events(events: List[Event]):
    """Watch for worker start/stop events and collect them"""
    async with get_events_subscriber(
        filter=EventFilter(
            event=EventNameFilter(prefix=["prefect.worker."]),
            occurred=EventOccurredFilter(since=now()),
        )
    ) as events_subscriber:
        async for event in events_subscriber:
            events.append(event)


def run_event_listener(events: List[Event]):
    """Run the async event listener in a thread"""
    asyncio.run(watch_worker_events(events))


def test_worker():
    WORKER_NAME = f"test-worker-{uuid4()}"  # noqa: F821
    events: List[Event] = []

    listener_thread = Thread(target=run_event_listener, args=(events,), daemon=True)
    listener_thread.start()

    try:
        subprocess.check_output(
            [
                uv.find_uv_bin(),
                "run",
                "--isolated",
                "prefect",
                "work-pool",
                "delete",
                "test-worker-pool",
            ],
        )
    except subprocess.CalledProcessError:
        pass

    try:
        subprocess.check_output(
            [
                uv.find_uv_bin(),
                "run",
                "--isolated",
                "--with",
                "prefect-kubernetes>=0.5.0",
                "prefect",
                "work-pool",
                "create",
                "test-worker-pool",
                "-t",
                "nonsense",
            ],
        )
    except subprocess.CalledProcessError as e:
        # Check that the error message contains kubernetes worker type
        for type in ["process", "kubernetes"]:
            assert type in str(e.output), (
                f"Worker type {type!r} missing from output {e.output}"
            )

    subprocess.check_call(
        [
            uv.find_uv_bin(),
            "run",
            "--isolated",
            "prefect",
            "work-pool",
            "create",
            "test-worker-pool",
            "-t",
            "kubernetes",
        ],
        stdout=sys.stdout,
        stderr=sys.stderr,
    )
    subprocess.check_call(
        [
            uv.find_uv_bin(),
            "run",
            "--isolated",
            "--with",
            "prefect-kubernetes",
            "prefect",
            "worker",
            "start",
            "-p",
            "test-worker-pool",
            "-t",
            "kubernetes",
            "-n",
            WORKER_NAME,
            "--run-once",
        ],
        stdout=sys.stdout,
        stderr=sys.stderr,
    )
    subprocess.check_call(
        [
            uv.find_uv_bin(),
            "run",
            "--isolated",
            "prefect",
            "--no-prompt",
            "work-pool",
            "delete",
            "test-worker-pool",
        ],
        stdout=sys.stdout,
        stderr=sys.stderr,
    )

    worker_events = [
        e
        for e in events
        if e.event.startswith("prefect.worker.") and e.resource.name == WORKER_NAME
    ]
    assert len(worker_events) == 2, (
        f"Expected 2 worker events, got {len(worker_events)}"
    )

    start_events = [e for e in worker_events if e.event == "prefect.worker.started"]
    stop_events = [e for e in worker_events if e.event == "prefect.worker.stopped"]

    assert len(start_events) == 1, "Expected 1 worker start event"
    assert len(stop_events) == 1, "Expected 1 worker stop event"

    print("Captured expected worker start and stop events!")

    assert stop_events[0].follows == start_events[0].id, (
        "Stop event should follow start event"
    )