port-ocean 0.28.5__py3-none-any.whl → 0.29.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (51) hide show
  1. integrations/_infra/Dockerfile.Deb +1 -0
  2. integrations/_infra/Dockerfile.local +1 -0
  3. port_ocean/clients/port/authentication.py +19 -0
  4. port_ocean/clients/port/client.py +3 -0
  5. port_ocean/clients/port/mixins/actions.py +93 -0
  6. port_ocean/clients/port/mixins/blueprints.py +0 -12
  7. port_ocean/clients/port/mixins/entities.py +79 -44
  8. port_ocean/clients/port/mixins/integrations.py +7 -2
  9. port_ocean/config/settings.py +35 -3
  10. port_ocean/context/ocean.py +7 -5
  11. port_ocean/core/defaults/initialize.py +12 -5
  12. port_ocean/core/event_listener/__init__.py +7 -0
  13. port_ocean/core/event_listener/actions_only.py +42 -0
  14. port_ocean/core/event_listener/base.py +4 -1
  15. port_ocean/core/event_listener/factory.py +18 -9
  16. port_ocean/core/event_listener/http.py +4 -3
  17. port_ocean/core/event_listener/kafka.py +3 -2
  18. port_ocean/core/event_listener/once.py +5 -2
  19. port_ocean/core/event_listener/polling.py +4 -3
  20. port_ocean/core/event_listener/webhooks_only.py +3 -2
  21. port_ocean/core/handlers/actions/__init__.py +7 -0
  22. port_ocean/core/handlers/actions/abstract_executor.py +150 -0
  23. port_ocean/core/handlers/actions/execution_manager.py +434 -0
  24. port_ocean/core/handlers/entity_processor/jq_entity_processor.py +479 -17
  25. port_ocean/core/handlers/entity_processor/jq_input_evaluator.py +137 -0
  26. port_ocean/core/handlers/port_app_config/models.py +4 -2
  27. port_ocean/core/handlers/webhook/abstract_webhook_processor.py +16 -0
  28. port_ocean/core/handlers/webhook/processor_manager.py +30 -12
  29. port_ocean/core/integrations/mixins/sync_raw.py +4 -4
  30. port_ocean/core/integrations/mixins/utils.py +250 -29
  31. port_ocean/core/models.py +35 -2
  32. port_ocean/core/utils/utils.py +16 -5
  33. port_ocean/exceptions/execution_manager.py +22 -0
  34. port_ocean/helpers/retry.py +4 -40
  35. port_ocean/log/logger_setup.py +2 -2
  36. port_ocean/ocean.py +30 -4
  37. port_ocean/tests/clients/port/mixins/test_entities.py +71 -5
  38. port_ocean/tests/core/event_listener/test_kafka.py +14 -7
  39. port_ocean/tests/core/handlers/actions/test_execution_manager.py +837 -0
  40. port_ocean/tests/core/handlers/entity_processor/test_jq_entity_processor.py +932 -1
  41. port_ocean/tests/core/handlers/entity_processor/test_jq_input_evaluator.py +932 -0
  42. port_ocean/tests/core/handlers/webhook/test_processor_manager.py +3 -1
  43. port_ocean/tests/core/utils/test_get_port_diff.py +164 -0
  44. port_ocean/tests/helpers/test_retry.py +241 -1
  45. port_ocean/tests/utils/test_cache.py +240 -0
  46. port_ocean/utils/cache.py +45 -9
  47. {port_ocean-0.28.5.dist-info → port_ocean-0.29.0.dist-info}/METADATA +2 -1
  48. {port_ocean-0.28.5.dist-info → port_ocean-0.29.0.dist-info}/RECORD +51 -41
  49. {port_ocean-0.28.5.dist-info → port_ocean-0.29.0.dist-info}/LICENSE.md +0 -0
  50. {port_ocean-0.28.5.dist-info → port_ocean-0.29.0.dist-info}/WHEEL +0 -0
  51. {port_ocean-0.28.5.dist-info → port_ocean-0.29.0.dist-info}/entry_points.txt +0 -0
@@ -0,0 +1,837 @@
1
+ import asyncio
2
+ from typing import Any
3
+ import uuid
4
+ from datetime import datetime, timedelta
5
+ from unittest.mock import ANY, AsyncMock, MagicMock, patch
6
+ from fastapi import APIRouter, FastAPI
7
+ from loguru import logger
8
+ from pydantic import BaseModel
9
+ import pytest
10
+ import httpx
11
+ from port_ocean.clients.port.authentication import PortAuthentication
12
+ from port_ocean.clients.port.client import PortClient
13
+ from port_ocean.context.ocean import PortOceanContext
14
+ from port_ocean.core.handlers.actions.abstract_executor import AbstractExecutor
15
+ from port_ocean.core.handlers.actions.execution_manager import (
16
+ ExecutionManager,
17
+ GLOBAL_SOURCE,
18
+ )
19
+ from port_ocean.core.handlers.webhook.abstract_webhook_processor import (
20
+ AbstractWebhookProcessor,
21
+ )
22
+ from port_ocean.core.handlers.webhook.processor_manager import (
23
+ LiveEventsProcessorManager,
24
+ )
25
+ from port_ocean.core.models import (
26
+ ActionRun,
27
+ IntegrationActionInvocationPayload,
28
+ IntegrationFeatureFlag,
29
+ RunStatus,
30
+ )
31
+ from port_ocean.exceptions.execution_manager import (
32
+ DuplicateActionExecutorError,
33
+ RunAlreadyAcknowledgedError,
34
+ )
35
+ from port_ocean.ocean import Ocean
36
+ from port_ocean.utils.signal import SignalHandler
37
+
38
+
39
+ def generate_mock_action_run(
40
+ action_type: str = "test_action",
41
+ integrationActionExecutionProperties: dict[str, Any] | None = None,
42
+ ) -> ActionRun:
43
+ if integrationActionExecutionProperties is None:
44
+ integrationActionExecutionProperties = {}
45
+ return ActionRun(
46
+ id=f"test-run-id-{uuid.uuid4()}",
47
+ status=RunStatus.IN_PROGRESS,
48
+ payload=IntegrationActionInvocationPayload(
49
+ type="INTEGRATION_ACTION",
50
+ installationId="test-installation-id",
51
+ integrationActionType=action_type,
52
+ integrationActionExecutionProperties=integrationActionExecutionProperties,
53
+ ),
54
+ )
55
+
56
+
57
+ @pytest.fixture
58
+ def mock_port_client() -> MagicMock:
59
+ mock_port_client = MagicMock(spec=PortClient)
60
+ mock_port_client.claim_pending_runs = AsyncMock()
61
+ mock_port_client.acknowledge_run = AsyncMock()
62
+ mock_port_client.get_run_by_external_id = AsyncMock()
63
+ mock_port_client.patch_run = AsyncMock()
64
+ mock_port_client.post_run_log = AsyncMock()
65
+ mock_port_client.get_organization_feature_flags = AsyncMock(
66
+ return_value=[IntegrationFeatureFlag.OCEAN_ACTIONS_PROCESSING_ENABLED]
67
+ )
68
+ mock_port_client.auth = AsyncMock(spec=PortAuthentication)
69
+ mock_port_client.auth.is_machine_user = AsyncMock(return_value=True)
70
+ return mock_port_client
71
+
72
+
73
+ @pytest.fixture
74
+ def mock_ocean(mock_port_client: PortClient) -> Ocean:
75
+ ocean_mock = MagicMock(spec=Ocean)
76
+ ocean_mock.config = MagicMock()
77
+ ocean_mock.port_client = mock_port_client
78
+ ocean_mock.integration_router = APIRouter()
79
+ ocean_mock.fast_api_app = FastAPI()
80
+ return ocean_mock
81
+
82
+
83
+ @pytest.fixture(autouse=True)
84
+ def mock_ocean_context(
85
+ monkeypatch: pytest.MonkeyPatch, mock_ocean: Ocean
86
+ ) -> PortOceanContext:
87
+ mock_ocean_context = PortOceanContext(mock_ocean)
88
+ mock_ocean_context._app = mock_ocean
89
+ monkeypatch.setattr(
90
+ "port_ocean.core.handlers.actions.execution_manager.ocean", mock_ocean_context
91
+ )
92
+ return mock_ocean_context
93
+
94
+
95
+ @pytest.fixture
96
+ def mock_webhook_manager() -> MagicMock:
97
+ return MagicMock(spec=LiveEventsProcessorManager)
98
+
99
+
100
+ @pytest.fixture
101
+ def mock_signal_handler() -> MagicMock:
102
+ return MagicMock(spec=SignalHandler)
103
+
104
+
105
+ @pytest.fixture
106
+ def mock_test_executor() -> MagicMock:
107
+ mock_executor = MagicMock(spec=AbstractExecutor)
108
+ mock_executor.ACTION_NAME = "test_action"
109
+ mock_executor.WEBHOOK_PROCESSOR_CLASS = None
110
+ mock_executor.WEBHOOK_PATH = None
111
+ mock_executor._get_partition_key = AsyncMock(return_value=None)
112
+ mock_executor.execute = AsyncMock(return_value=None)
113
+ mock_executor.is_close_to_rate_limit = AsyncMock(return_value=False)
114
+ mock_executor.get_remaining_seconds_until_rate_limit = AsyncMock(return_value=0.0)
115
+ return mock_executor
116
+
117
+
118
+ @pytest.fixture
119
+ def mock_test_partition_executor() -> MagicMock:
120
+ mock_executor = MagicMock(spec=AbstractExecutor)
121
+ mock_executor.ACTION_NAME = "test_partition_action"
122
+ mock_executor.WEBHOOK_PROCESSOR_CLASS = None
123
+ mock_executor.WEBHOOK_PATH = None
124
+ mock_executor._get_partition_key = AsyncMock(
125
+ side_effect=lambda run: run.payload.integrationActionExecutionProperties.get(
126
+ "partition_name", "default_partition"
127
+ )
128
+ )
129
+ mock_executor.execute = AsyncMock(return_value=None)
130
+ mock_executor.is_close_to_rate_limit = AsyncMock(return_value=False)
131
+ mock_executor.get_remaining_seconds_until_rate_limit = AsyncMock(return_value=0.0)
132
+ return mock_executor
133
+
134
+
135
+ @pytest.fixture
136
+ def execution_manager_without_executors(
137
+ mock_webhook_manager: MagicMock, mock_signal_handler: MagicMock
138
+ ) -> ExecutionManager:
139
+ return ExecutionManager(
140
+ webhook_manager=mock_webhook_manager,
141
+ signal_handler=mock_signal_handler,
142
+ workers_count=3,
143
+ runs_buffer_high_watermark=100,
144
+ poll_check_interval_seconds=5,
145
+ visibility_timeout_ms=30,
146
+ max_wait_seconds_before_shutdown=30,
147
+ )
148
+
149
+
150
+ @pytest.fixture
151
+ def execution_manager(
152
+ mock_webhook_manager: MagicMock,
153
+ mock_signal_handler: MagicMock,
154
+ mock_test_executor: MagicMock,
155
+ mock_test_partition_executor: MagicMock,
156
+ ) -> ExecutionManager:
157
+ execution_manager = ExecutionManager(
158
+ webhook_manager=mock_webhook_manager,
159
+ signal_handler=mock_signal_handler,
160
+ workers_count=3,
161
+ runs_buffer_high_watermark=100,
162
+ poll_check_interval_seconds=5,
163
+ visibility_timeout_ms=30,
164
+ max_wait_seconds_before_shutdown=30,
165
+ )
166
+ execution_manager.register_executor(mock_test_executor)
167
+ execution_manager.register_executor(mock_test_partition_executor)
168
+ return execution_manager
169
+
170
+
171
+ class TestExecutionManager:
172
+ @pytest.mark.asyncio
173
+ async def test_register_executor(
174
+ self,
175
+ execution_manager_without_executors: ExecutionManager,
176
+ mock_webhook_manager: MagicMock,
177
+ mock_test_executor: MagicMock,
178
+ ) -> None:
179
+ # Arrange
180
+ mock_test_executor.WEBHOOK_PROCESSOR_CLASS = MagicMock(
181
+ spec=AbstractWebhookProcessor
182
+ )
183
+ mock_test_executor.WEBHOOK_PATH = "/test-webhook"
184
+
185
+ # Act
186
+ execution_manager_without_executors.register_executor(mock_test_executor)
187
+
188
+ # Assert
189
+ assert (
190
+ mock_test_executor.ACTION_NAME
191
+ in execution_manager_without_executors._actions_executors
192
+ )
193
+ mock_webhook_manager.register_processor.assert_called_once_with(
194
+ mock_test_executor.WEBHOOK_PATH,
195
+ mock_test_executor.WEBHOOK_PROCESSOR_CLASS,
196
+ )
197
+
198
+ @pytest.mark.asyncio
199
+ async def test_register_executor_should_raise_error_if_duplicate(
200
+ self, execution_manager: ExecutionManager, mock_test_executor: MagicMock
201
+ ) -> None:
202
+ # Act & Assert
203
+ with pytest.raises(
204
+ DuplicateActionExecutorError,
205
+ match="Executor for action 'test_action' is already registered",
206
+ ):
207
+ execution_manager.register_executor(mock_test_executor)
208
+
209
+ @pytest.mark.asyncio
210
+ async def test_execute_run_should_acknowledge_run_successfully(
211
+ self,
212
+ execution_manager: ExecutionManager,
213
+ mock_port_client: MagicMock,
214
+ ) -> None:
215
+ # Arrange
216
+ mock_test_action_run = generate_mock_action_run()
217
+
218
+ # Act & Assert
219
+ with patch.object(
220
+ execution_manager._actions_executors["test_action"], "execute"
221
+ ) as mock_execute:
222
+ await execution_manager._execute_run(mock_test_action_run)
223
+ mock_port_client.acknowledge_run.assert_called_once_with(
224
+ mock_test_action_run.id
225
+ )
226
+ mock_execute.assert_called_once()
227
+
228
+ @pytest.mark.asyncio
229
+ async def test_execute_run_should_not_execute_run_if_acknowledge_conflicts(
230
+ self,
231
+ execution_manager: ExecutionManager,
232
+ mock_port_client: MagicMock,
233
+ ) -> None:
234
+ # Arrange
235
+ mock_test_action_run = generate_mock_action_run()
236
+ mock_port_client.acknowledge_run.side_effect = RunAlreadyAcknowledgedError()
237
+
238
+ # Act
239
+ with patch.object(
240
+ execution_manager._actions_executors["test_action"], "execute"
241
+ ) as mock_execute:
242
+ await execution_manager._execute_run(mock_test_action_run)
243
+
244
+ # Assert
245
+ mock_port_client.acknowledge_run.assert_called_once_with(
246
+ mock_test_action_run.id
247
+ )
248
+ mock_execute.assert_not_called()
249
+
250
+ @pytest.mark.asyncio
251
+ async def test_execute_run_should_sleep_if_rate_limited(
252
+ self,
253
+ execution_manager_without_executors: ExecutionManager,
254
+ mock_port_client: MagicMock,
255
+ mock_test_executor: MagicMock,
256
+ ) -> None:
257
+ # Arrange
258
+ few_seconds_away = datetime.now() + timedelta(seconds=0.1)
259
+ mock_test_executor.is_close_to_rate_limit = AsyncMock(
260
+ side_effect=lambda: few_seconds_away > datetime.now()
261
+ )
262
+ mock_test_executor.get_remaining_seconds_until_rate_limit = AsyncMock(
263
+ side_effect=lambda: (few_seconds_away - datetime.now()).total_seconds()
264
+ )
265
+ execution_manager_without_executors.register_executor(mock_test_executor)
266
+ mock_test_action_run = generate_mock_action_run()
267
+
268
+ # Act
269
+ await execution_manager_without_executors._execute_run(mock_test_action_run)
270
+
271
+ # Assert
272
+ mock_port_client.post_run_log.assert_called_with(
273
+ mock_test_action_run.id,
274
+ ANY,
275
+ )
276
+
277
+ @pytest.mark.asyncio
278
+ async def test_get_queues_size_with_global_and_partition_queues(
279
+ self, execution_manager: ExecutionManager
280
+ ) -> None:
281
+ # Arrange
282
+ await execution_manager._global_queue.put(generate_mock_action_run())
283
+ await execution_manager._global_queue.put(generate_mock_action_run())
284
+
285
+ execution_manager._partition_queues["partition1"] = (
286
+ execution_manager._global_queue.__class__()
287
+ )
288
+ await execution_manager._partition_queues["partition1"].put(
289
+ generate_mock_action_run()
290
+ )
291
+
292
+ # Act
293
+ size = await execution_manager._get_queues_size()
294
+
295
+ # Assert
296
+ assert size == 3
297
+
298
+ @pytest.mark.asyncio
299
+ async def test_add_run_to_queue_should_add_source_to_active_when_empty(
300
+ self, execution_manager: ExecutionManager
301
+ ) -> None:
302
+ # Arrange
303
+ run = generate_mock_action_run()
304
+
305
+ # Act
306
+ await execution_manager._add_run_to_queue(run, GLOBAL_SOURCE)
307
+
308
+ # Assert
309
+ assert await execution_manager._global_queue.size() == 1
310
+ active_source = await execution_manager._active_sources.get()
311
+ assert active_source == GLOBAL_SOURCE
312
+ assert run.id in execution_manager._deduplication_set
313
+
314
+ @pytest.mark.asyncio
315
+ async def test_add_run_to_queue_should_not_add_source_when_queue_has_items(
316
+ self, execution_manager: ExecutionManager
317
+ ) -> None:
318
+ # Arrange
319
+ run1 = generate_mock_action_run()
320
+ run2 = generate_mock_action_run()
321
+
322
+ # Act & Assert
323
+ await execution_manager._add_run_to_queue(run1, GLOBAL_SOURCE)
324
+
325
+ with patch.object(execution_manager._active_sources, "put") as mock_add_source:
326
+ await execution_manager._add_run_to_queue(run2, GLOBAL_SOURCE)
327
+ mock_add_source.assert_not_called()
328
+
329
+ @pytest.mark.asyncio
330
+ async def test_add_run_to_queue_should_create_queue_if_not_exists(
331
+ self, execution_manager: ExecutionManager
332
+ ) -> None:
333
+ # Arrange
334
+ queue_name = "test_action:partition1"
335
+ run = generate_mock_action_run()
336
+
337
+ # Act
338
+ await execution_manager._add_run_to_queue(run, queue_name)
339
+
340
+ # Assert
341
+ assert queue_name in execution_manager._partition_queues
342
+ assert queue_name in execution_manager._queues_locks
343
+ assert await execution_manager._partition_queues[queue_name].size() == 1
344
+
345
+ @pytest.mark.asyncio
346
+ async def test_handle_global_queue_once_should_process_run_and_remove_dedup(
347
+ self,
348
+ execution_manager: ExecutionManager,
349
+ mock_port_client: MagicMock,
350
+ ) -> None:
351
+ # Arrange
352
+ run = generate_mock_action_run()
353
+ await execution_manager._add_run_to_queue(
354
+ run,
355
+ GLOBAL_SOURCE,
356
+ )
357
+
358
+ # Act & Assert
359
+ with patch.object(
360
+ execution_manager._actions_executors["test_action"], "execute"
361
+ ) as mock_execute:
362
+ await execution_manager._handle_global_queue_once()
363
+
364
+ assert run.id not in execution_manager._deduplication_set
365
+ mock_port_client.acknowledge_run.assert_called_once_with(run.id)
366
+ mock_execute.assert_called_once_with(run)
367
+
368
+ @pytest.mark.asyncio
369
+ async def test_handle_partition_queue_once_should_process_run(
370
+ self,
371
+ execution_manager: ExecutionManager,
372
+ mock_port_client: MagicMock,
373
+ ) -> None:
374
+ # Arrange
375
+ partition_name = "test_action:partition1"
376
+ run = generate_mock_action_run()
377
+ await execution_manager._add_run_to_queue(
378
+ run,
379
+ partition_name,
380
+ )
381
+
382
+ # Act & Assert
383
+ with patch.object(
384
+ execution_manager._actions_executors["test_action"], "execute"
385
+ ) as mock_execute:
386
+ await execution_manager._handle_partition_queue_once(partition_name)
387
+
388
+ assert run.id not in execution_manager._deduplication_set
389
+ mock_port_client.acknowledge_run.assert_called_once_with(run.id)
390
+ mock_execute.assert_called_once_with(run)
391
+
392
+ @pytest.mark.asyncio
393
+ async def test_poll_action_runs_should_respect_high_watermark(
394
+ self, execution_manager: ExecutionManager, mock_port_client: MagicMock
395
+ ) -> None:
396
+ # Arrange
397
+ execution_manager._high_watermark = 2
398
+ for _ in range(3):
399
+ await execution_manager._global_queue.put(generate_mock_action_run())
400
+
401
+ mock_port_client.claim_pending_runs.return_value = []
402
+
403
+ # Act
404
+ polling_task: asyncio.Task[None] = asyncio.create_task(
405
+ execution_manager._poll_action_runs()
406
+ )
407
+ await asyncio.sleep(0.1)
408
+ await execution_manager._gracefully_cancel_task(polling_task)
409
+
410
+ # Assert
411
+ mock_port_client.claim_pending_runs.assert_not_called()
412
+
413
+ @pytest.mark.asyncio
414
+ async def test_poll_action_runs_should_poll_when_below_watermark(
415
+ self,
416
+ execution_manager: ExecutionManager,
417
+ mock_port_client: MagicMock,
418
+ ) -> None:
419
+ # Arrange
420
+ execution_manager._high_watermark = 10
421
+ execution_manager._poll_check_interval_seconds = 0
422
+ mock_port_client.claim_pending_runs.side_effect = (
423
+ lambda limit, visibility_timeout_ms: [generate_mock_action_run()]
424
+ )
425
+
426
+ # Act
427
+ polling_task: asyncio.Task[None] = asyncio.create_task(
428
+ execution_manager._poll_action_runs()
429
+ )
430
+ await asyncio.sleep(0.1)
431
+ await execution_manager._gracefully_cancel_task(polling_task)
432
+
433
+ # Assert
434
+ mock_port_client.claim_pending_runs.assert_called()
435
+ assert (
436
+ await execution_manager._global_queue.size()
437
+ == execution_manager._high_watermark
438
+ )
439
+
440
+ @pytest.mark.asyncio
441
+ async def test_poll_action_runs_should_skip_unregistered_actions(
442
+ self, execution_manager: ExecutionManager, mock_port_client: MagicMock
443
+ ) -> None:
444
+ # Arrange
445
+ execution_manager._high_watermark = 10
446
+ execution_manager._poll_check_interval_seconds = 0
447
+ mock_port_client.claim_pending_runs.side_effect = (
448
+ lambda limit, visibility_timeout_ms: [
449
+ generate_mock_action_run(action_type="unregistered_action")
450
+ ]
451
+ )
452
+
453
+ # Act
454
+ polling_task: asyncio.Task[None] = asyncio.create_task(
455
+ execution_manager._poll_action_runs()
456
+ )
457
+ await asyncio.sleep(0.1)
458
+ await execution_manager._gracefully_cancel_task(polling_task)
459
+
460
+ # Assert
461
+ assert await execution_manager._get_queues_size() == 0
462
+
463
+ @pytest.mark.asyncio
464
+ async def test_shutdown_should_cancel_polling_and_waits_for_workers(
465
+ self,
466
+ execution_manager: ExecutionManager,
467
+ ) -> None:
468
+ # Arrange
469
+ execution_manager._max_wait_seconds_before_shutdown = 1.0
470
+ await execution_manager.start_processing_action_runs()
471
+
472
+ # Act
473
+ assert execution_manager._polling_task is not None
474
+ assert (
475
+ not execution_manager._polling_task.done()
476
+ or execution_manager._polling_task.cancelled()
477
+ )
478
+ assert not execution_manager._is_shutting_down.is_set()
479
+ assert len(execution_manager._workers_pool) == 3
480
+ await execution_manager.shutdown()
481
+
482
+ # Assert
483
+ assert execution_manager._is_shutting_down.is_set()
484
+ assert execution_manager._polling_task.cancelled()
485
+ assert len(execution_manager._workers_pool) == 0
486
+
487
+ @pytest.mark.asyncio
488
+ async def test_shutdown_should_not_acknowledge_runs_after_shutdown_started(
489
+ self,
490
+ execution_manager: ExecutionManager,
491
+ mock_port_client: MagicMock,
492
+ ) -> None:
493
+ # Arrange
494
+ execution_manager._max_wait_seconds_before_shutdown = 1.0
495
+ execution_manager._high_watermark = 50
496
+ execution_manager._poll_check_interval_seconds = 0
497
+ mock_port_client.claim_pending_runs.return_value = [
498
+ generate_mock_action_run() for _ in range(10)
499
+ ]
500
+ await execution_manager.start_processing_action_runs()
501
+
502
+ # Act
503
+ ack_calls_count: int = mock_port_client.acknowledge_run.call_count
504
+ await execution_manager.shutdown()
505
+
506
+ # Assert
507
+ assert mock_port_client.acknowledge_run.call_count == ack_calls_count
508
+
509
+ @pytest.mark.asyncio
510
+ async def test_global_and_partition_queues_concurrency(
511
+ self,
512
+ execution_manager_without_executors: ExecutionManager,
513
+ mock_test_executor: MagicMock,
514
+ mock_test_partition_executor: MagicMock,
515
+ mock_port_client: MagicMock,
516
+ ) -> None:
517
+ """Test that partition queues process sequentially while global queue processes concurrently"""
518
+
519
+ # Arrange
520
+ class RunMeasurement(BaseModel):
521
+ start_time: datetime
522
+ end_time: datetime
523
+
524
+ partition1 = "partition1"
525
+ partition2 = "partition2"
526
+ partition1_queue_name = f"{mock_test_partition_executor.ACTION_NAME}:partition1"
527
+ partition2_queue_name = f"{mock_test_partition_executor.ACTION_NAME}:partition2"
528
+ run_measurements: dict[str, list[RunMeasurement]] = {
529
+ partition1_queue_name: [],
530
+ partition2_queue_name: [],
531
+ GLOBAL_SOURCE: [],
532
+ }
533
+
534
+ execution_manager_without_executors._workers_count = 5
535
+ execution_manager_without_executors._high_watermark = 20
536
+ execution_manager_without_executors._poll_check_interval_seconds = 0
537
+ execution_manager_without_executors._max_wait_seconds_before_shutdown = 1.0
538
+
539
+ async def mock_execute(
540
+ run: ActionRun,
541
+ ) -> None:
542
+ await asyncio.sleep(0.1)
543
+ return None
544
+
545
+ mock_test_executor.execute.side_effect = mock_execute
546
+ mock_test_partition_executor.execute.side_effect = mock_execute
547
+ execution_manager_without_executors.register_executor(mock_test_executor)
548
+ execution_manager_without_executors.register_executor(
549
+ mock_test_partition_executor
550
+ )
551
+
552
+ # Patch the relevant methods to record execution timings for measurement
553
+ original_handle_global_queue_once = (
554
+ execution_manager_without_executors._handle_global_queue_once
555
+ )
556
+ original_handle_partition_queue_once = (
557
+ execution_manager_without_executors._handle_partition_queue_once
558
+ )
559
+
560
+ async def wrapped_handle_global_queue_once() -> None:
561
+ start_time = datetime.now()
562
+ await original_handle_global_queue_once()
563
+ try:
564
+ run_measurements[GLOBAL_SOURCE].append(
565
+ RunMeasurement(start_time=start_time, end_time=datetime.now())
566
+ )
567
+ except Exception as e:
568
+ logger.error(f"Error recording run measurement: {e}")
569
+
570
+ async def wrapped_handle_partition_queue_once(partition_name: str) -> None:
571
+ start_time = datetime.now()
572
+ await original_handle_partition_queue_once(partition_name)
573
+ try:
574
+ run_measurements[partition_name].append(
575
+ RunMeasurement(start_time=start_time, end_time=datetime.now())
576
+ )
577
+ except Exception as e:
578
+ logger.error(f"Error recording run measurement: {e}")
579
+
580
+ setattr(
581
+ execution_manager_without_executors,
582
+ "_handle_global_queue_once",
583
+ wrapped_handle_global_queue_once,
584
+ )
585
+ setattr(
586
+ execution_manager_without_executors,
587
+ "_handle_partition_queue_once",
588
+ wrapped_handle_partition_queue_once,
589
+ )
590
+ mock_port_client.claim_pending_runs.side_effect = (
591
+ lambda limit, visibility_timeout_ms: [
592
+ *[
593
+ generate_mock_action_run(
594
+ action_type=mock_test_partition_executor.ACTION_NAME,
595
+ integrationActionExecutionProperties={
596
+ "partition_name": partition1
597
+ },
598
+ )
599
+ for _ in range(5)
600
+ ],
601
+ *[
602
+ generate_mock_action_run(
603
+ action_type=mock_test_partition_executor.ACTION_NAME,
604
+ integrationActionExecutionProperties={
605
+ "partition_name": partition2
606
+ },
607
+ )
608
+ for _ in range(5)
609
+ ],
610
+ *[generate_mock_action_run() for _ in range(5)],
611
+ ]
612
+ )
613
+
614
+ def check_global_queue_measurements() -> None:
615
+ queue_measurements = run_measurements[GLOBAL_SOURCE]
616
+ assert any(
617
+ m.end_time >= queue_measurements[i + 1].start_time
618
+ for i, m in enumerate(queue_measurements[:-1])
619
+ )
620
+
621
+ def check_partition_queue_measurements(queue_name: str) -> None:
622
+ queue_measurements = run_measurements[queue_name]
623
+ for idx, measurement in enumerate(queue_measurements):
624
+ if idx == len(queue_measurements) - 1:
625
+ continue
626
+ assert measurement.end_time < queue_measurements[idx + 1].start_time
627
+
628
+ # Act
629
+ await execution_manager_without_executors.start_processing_action_runs()
630
+ await asyncio.sleep(1)
631
+ await execution_manager_without_executors.shutdown()
632
+
633
+ # Assert
634
+ assert execution_manager_without_executors._polling_task is not None
635
+ assert execution_manager_without_executors._polling_task.cancelled()
636
+ assert len(execution_manager_without_executors._workers_pool) == 0
637
+ for queue_name in run_measurements:
638
+ assert len(run_measurements[queue_name]) > 0
639
+ if queue_name == GLOBAL_SOURCE:
640
+ check_global_queue_measurements()
641
+ else:
642
+ check_partition_queue_measurements(queue_name)
643
+
644
+ @pytest.mark.asyncio
645
+ async def test_execute_run_handles_general_exception(
646
+ self,
647
+ execution_manager: ExecutionManager,
648
+ mock_port_client: MagicMock,
649
+ mock_test_executor: MagicMock,
650
+ ) -> None:
651
+ # Arrange
652
+ run = generate_mock_action_run()
653
+ error_msg = "Test error"
654
+ mock_test_executor.execute.side_effect = Exception(error_msg)
655
+
656
+ # Act
657
+ await execution_manager._execute_run(run)
658
+
659
+ # Assert
660
+ assert mock_port_client.patch_run.call_count == 1
661
+ called_args, _ = mock_port_client.patch_run.call_args
662
+ assert called_args[0] == run.id
663
+ patch_data = called_args[1]
664
+ assert error_msg in patch_data["summary"]
665
+ assert patch_data["status"] == RunStatus.FAILURE
666
+
667
+ @pytest.mark.asyncio
668
+ async def test_execute_run_handles_acknowledge_run_api_error(
669
+ self,
670
+ execution_manager: ExecutionManager,
671
+ mock_port_client: MagicMock,
672
+ ) -> None:
673
+ # Arrange
674
+ run = generate_mock_action_run()
675
+ mock_response = MagicMock(spec=httpx.Response)
676
+ mock_response.status_code = 500
677
+ mock_response.text = "Internal Server Error"
678
+ http_error = httpx.HTTPStatusError(
679
+ "500 Internal Server Error",
680
+ request=MagicMock(),
681
+ response=mock_response,
682
+ )
683
+ mock_port_client.acknowledge_run.side_effect = http_error
684
+
685
+ # Act
686
+ await execution_manager._execute_run(run)
687
+
688
+ # Assert
689
+ assert mock_port_client.patch_run.call_count == 1
690
+ called_args, _ = mock_port_client.patch_run.call_args
691
+ assert called_args[0] == run.id
692
+ patch_data = called_args[1]
693
+ patch_data["summary"] == "Failed to trigger run execution"
694
+ assert patch_data["status"] == RunStatus.FAILURE
695
+
696
+ @pytest.mark.asyncio
697
+ async def test_polling_continues_after_api_errors(
698
+ self,
699
+ execution_manager: ExecutionManager,
700
+ mock_port_client: MagicMock,
701
+ ) -> None:
702
+ """Verify that polling loop continues running even after multiple API errors"""
703
+ # Arrange
704
+ execution_manager._high_watermark = 10
705
+ execution_manager._poll_check_interval_seconds = 0
706
+
707
+ poll_count = 0
708
+
709
+ async def claim_runs_with_errors(
710
+ limit: int, visibility_timeout_ms: int
711
+ ) -> list[ActionRun]:
712
+ nonlocal poll_count
713
+ poll_count += 1
714
+ # Fail on first and third attempts, succeed on second and fourth
715
+ if poll_count in [1, 3]:
716
+ mock_response = MagicMock(spec=httpx.Response)
717
+ mock_response.status_code = 500
718
+ raise httpx.HTTPStatusError(
719
+ "500 Internal Server Error",
720
+ request=MagicMock(),
721
+ response=mock_response,
722
+ )
723
+ return [generate_mock_action_run()]
724
+
725
+ mock_port_client.claim_pending_runs.side_effect = claim_runs_with_errors
726
+
727
+ # Act
728
+ polling_task: asyncio.Task[None] = asyncio.create_task(
729
+ execution_manager._poll_action_runs()
730
+ )
731
+ await asyncio.sleep(0.2)
732
+ await execution_manager._gracefully_cancel_task(polling_task)
733
+
734
+ # Assert
735
+ # Polling should have continued through errors
736
+ assert poll_count >= 4
737
+ assert mock_port_client.claim_pending_runs.call_count >= 4
738
+ # Should have successfully processed runs from successful polls
739
+ assert await execution_manager._get_queues_size() > 0
740
+
741
+ @pytest.mark.asyncio
742
+ async def test_process_actions_runs_handles_exceptions_gracefully(
743
+ self,
744
+ execution_manager: ExecutionManager,
745
+ mock_port_client: MagicMock,
746
+ mock_test_executor: MagicMock,
747
+ ) -> None:
748
+ """Verify that _process_actions_runs catches exceptions and doesn't crash the worker"""
749
+ # Arrange
750
+ run = generate_mock_action_run()
751
+
752
+ # Make execute raise an exception, and patch_run also fail
753
+ mock_test_executor.execute.side_effect = Exception("Execution failed")
754
+ mock_response = MagicMock(spec=httpx.Response)
755
+ mock_response.status_code = 503
756
+ patch_error = httpx.HTTPStatusError(
757
+ "503 Service Unavailable",
758
+ request=MagicMock(),
759
+ response=mock_response,
760
+ )
761
+ mock_port_client.patch_run.side_effect = patch_error
762
+
763
+ # Add run to queue
764
+ await execution_manager._add_run_to_queue(run, GLOBAL_SOURCE)
765
+
766
+ # Act
767
+ # Start worker loop which should handle exceptions gracefully
768
+ worker_task = asyncio.create_task(execution_manager._process_actions_runs())
769
+
770
+ # Wait for worker to process the run
771
+ await asyncio.sleep(0.1)
772
+
773
+ # Signal shutdown - worker should still be running (not crashed)
774
+ execution_manager._is_shutting_down.set()
775
+
776
+ # Wait a bit more to ensure worker handles shutdown
777
+ await asyncio.sleep(0.1)
778
+
779
+ # Cancel worker
780
+ worker_task.cancel()
781
+ try:
782
+ await worker_task
783
+ except asyncio.CancelledError:
784
+ pass
785
+
786
+ # Assert
787
+ # Exception should have been caught by worker loop, not crashed
788
+ # Run should have been acknowledged before execution failed
789
+ mock_port_client.acknowledge_run.assert_called_once_with(run.id)
790
+ # patch_run should have been attempted (even if it failed)
791
+ mock_port_client.patch_run.assert_called_once()
792
+ # Worker task should have completed (either naturally or cancelled), not crashed
793
+ assert worker_task.done()
794
+
795
+ @pytest.mark.asyncio
796
+ async def test_poll_action_runs_continues_after_error_handling_runs(
797
+ self,
798
+ execution_manager: ExecutionManager,
799
+ mock_port_client: MagicMock,
800
+ ) -> None:
801
+ # Arrange
802
+ execution_manager._high_watermark = 10
803
+ execution_manager._poll_check_interval_seconds = 0
804
+
805
+ # First call succeeds, second call fails, third call succeeds again
806
+ call_count = 0
807
+
808
+ async def claim_runs_side_effect(
809
+ limit: int, visibility_timeout_ms: int
810
+ ) -> list[ActionRun]:
811
+ nonlocal call_count
812
+ call_count += 1
813
+ if call_count == 2:
814
+ mock_response = MagicMock(spec=httpx.Response)
815
+ mock_response.status_code = 500
816
+ mock_response.text = "Internal Server Error"
817
+ raise httpx.HTTPStatusError(
818
+ "500 Internal Server Error",
819
+ request=MagicMock(),
820
+ response=mock_response,
821
+ )
822
+ return [generate_mock_action_run()]
823
+
824
+ mock_port_client.claim_pending_runs.side_effect = claim_runs_side_effect
825
+
826
+ # Act
827
+ polling_task: asyncio.Task[None] = asyncio.create_task(
828
+ execution_manager._poll_action_runs()
829
+ )
830
+ await asyncio.sleep(0.15) # Allow multiple poll attempts
831
+ await execution_manager._gracefully_cancel_task(polling_task)
832
+
833
+ # Assert
834
+ # Should have attempted to poll multiple times, handling the error gracefully
835
+ assert mock_port_client.claim_pending_runs.call_count >= 2
836
+ # Should have successfully added runs from successful polls
837
+ assert await execution_manager._get_queues_size() > 0