avtomatika 1.0b1__tar.gz → 1.0b2__tar.gz
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.
- {avtomatika-1.0b1/src/avtomatika.egg-info → avtomatika-1.0b2}/PKG-INFO +9 -2
- {avtomatika-1.0b1 → avtomatika-1.0b2}/README.md +8 -1
- {avtomatika-1.0b1 → avtomatika-1.0b2}/pyproject.toml +1 -1
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/api.html +14 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/blueprint.py +8 -1
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/engine.py +13 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/security.py +5 -3
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/worker_config_loader.py +4 -1
- {avtomatika-1.0b1 → avtomatika-1.0b2/src/avtomatika.egg-info}/PKG-INFO +9 -2
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika.egg-info/SOURCES.txt +1 -0
- avtomatika-1.0b2/tests/test_compression.py +121 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_integration.py +9 -2
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_worker_config_loader.py +6 -2
- {avtomatika-1.0b1 → avtomatika-1.0b2}/LICENSE +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/setup.cfg +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/__init__.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/client_config_loader.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/compression.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/config.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/context.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/data_types.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/datastore.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/dispatcher.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/executor.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/health_checker.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/history/base.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/history/noop.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/history/postgres.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/history/sqlite.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/logging_config.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/metrics.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/py.typed +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/quota.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/ratelimit.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/reputation.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/storage/__init__.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/storage/base.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/storage/memory.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/storage/redis.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/telemetry.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/watcher.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika/ws_manager.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika.egg-info/dependency_links.txt +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika.egg-info/requires.txt +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/src/avtomatika.egg-info/top_level.txt +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_blueprint_conditions.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_blueprints.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_client_config_loader.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_context.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_dispatcher.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_engine.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_error_handling.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_executor.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_health_checker.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_history.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_logging_config.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_memory_storage.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_metrics.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_noop_history.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_postgres_history.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_ratelimit.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_redis_storage.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_reputation.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_telemetry.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_watcher.py +0 -0
- {avtomatika-1.0b1 → avtomatika-1.0b2}/tests/test_ws_manager.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: avtomatika
|
|
3
|
-
Version: 1.
|
|
3
|
+
Version: 1.0b2
|
|
4
4
|
Summary: A state-machine based orchestrator for long-running jobs.
|
|
5
5
|
Project-URL: Homepage, https://github.com/avtomatika-ai/avtomatika
|
|
6
6
|
Project-URL: Bug Tracker, https://github.com/avtomatika-ai/avtomatika/issues
|
|
@@ -251,6 +251,11 @@ async def handle_normal(actions):
|
|
|
251
251
|
actions.transition_to("normal_processing")
|
|
252
252
|
```
|
|
253
253
|
|
|
254
|
+
> **Note on Limitations:** The current version of `.when()` uses a simple parser with the following limitations:
|
|
255
|
+
> * **No Nested Attributes:** You can only access direct fields of `context.initial_data` or `context.state_history` (e.g., `context.initial_data.field`). Nested objects (e.g., `context.initial_data.area.field`) are not supported.
|
|
256
|
+
> * **Simple Comparisons Only:** Only the following operators are supported: `==`, `!=`, `>`, `<`, `>=`, `<=`. Complex logical expressions with `AND`, `OR`, or `NOT` are not allowed.
|
|
257
|
+
> * **Limited Value Types:** The parser only recognizes strings (in quotes), integers, and floats. Boolean values (`True`, `False`) and `None` are not correctly parsed and will be treated as strings.
|
|
258
|
+
|
|
254
259
|
### 2. Delegating Tasks to Workers (`dispatch_task`)
|
|
255
260
|
|
|
256
261
|
This is the primary function for delegating work. The orchestrator will queue the task and wait for a worker to pick it up and return a result.
|
|
@@ -368,7 +373,9 @@ The orchestrator uses tokens to authenticate API requests.
|
|
|
368
373
|
* **Client Authentication**: All API clients must provide a token in the `X-Avtomatika-Token` header. The orchestrator validates this token against client configurations.
|
|
369
374
|
* **Worker Authentication**: Workers must provide a token in the `X-Worker-Token` header.
|
|
370
375
|
* `GLOBAL_WORKER_TOKEN`: You can set a global token for all workers using this environment variable. For development and testing, it defaults to `"secure-worker-token"`.
|
|
371
|
-
* **Individual Tokens**: For production, it is recommended to define individual tokens for each worker in a separate configuration file and provide its path via the `WORKERS_CONFIG_PATH` environment variable.
|
|
376
|
+
* **Individual Tokens**: For production, it is recommended to define individual tokens for each worker in a separate configuration file and provide its path via the `WORKERS_CONFIG_PATH` environment variable. Tokens from this file are stored in a hashed format for security.
|
|
377
|
+
|
|
378
|
+
> **Note on Dynamic Reloading:** The worker configuration file can be reloaded without restarting the orchestrator by sending an authenticated `POST` request to the `/api/v1/admin/reload-workers` endpoint. This allows for dynamic updates of worker tokens.
|
|
372
379
|
|
|
373
380
|
### Observability
|
|
374
381
|
|
|
@@ -205,6 +205,11 @@ async def handle_normal(actions):
|
|
|
205
205
|
actions.transition_to("normal_processing")
|
|
206
206
|
```
|
|
207
207
|
|
|
208
|
+
> **Note on Limitations:** The current version of `.when()` uses a simple parser with the following limitations:
|
|
209
|
+
> * **No Nested Attributes:** You can only access direct fields of `context.initial_data` or `context.state_history` (e.g., `context.initial_data.field`). Nested objects (e.g., `context.initial_data.area.field`) are not supported.
|
|
210
|
+
> * **Simple Comparisons Only:** Only the following operators are supported: `==`, `!=`, `>`, `<`, `>=`, `<=`. Complex logical expressions with `AND`, `OR`, or `NOT` are not allowed.
|
|
211
|
+
> * **Limited Value Types:** The parser only recognizes strings (in quotes), integers, and floats. Boolean values (`True`, `False`) and `None` are not correctly parsed and will be treated as strings.
|
|
212
|
+
|
|
208
213
|
### 2. Delegating Tasks to Workers (`dispatch_task`)
|
|
209
214
|
|
|
210
215
|
This is the primary function for delegating work. The orchestrator will queue the task and wait for a worker to pick it up and return a result.
|
|
@@ -322,7 +327,9 @@ The orchestrator uses tokens to authenticate API requests.
|
|
|
322
327
|
* **Client Authentication**: All API clients must provide a token in the `X-Avtomatika-Token` header. The orchestrator validates this token against client configurations.
|
|
323
328
|
* **Worker Authentication**: Workers must provide a token in the `X-Worker-Token` header.
|
|
324
329
|
* `GLOBAL_WORKER_TOKEN`: You can set a global token for all workers using this environment variable. For development and testing, it defaults to `"secure-worker-token"`.
|
|
325
|
-
* **Individual Tokens**: For production, it is recommended to define individual tokens for each worker in a separate configuration file and provide its path via the `WORKERS_CONFIG_PATH` environment variable.
|
|
330
|
+
* **Individual Tokens**: For production, it is recommended to define individual tokens for each worker in a separate configuration file and provide its path via the `WORKERS_CONFIG_PATH` environment variable. Tokens from this file are stored in a hashed format for security.
|
|
331
|
+
|
|
332
|
+
> **Note on Dynamic Reloading:** The worker configuration file can be reloaded without restarting the orchestrator by sending an authenticated `POST` request to the `/api/v1/admin/reload-workers` endpoint. This allows for dynamic updates of worker tokens.
|
|
326
333
|
|
|
327
334
|
### Observability
|
|
328
335
|
|
|
@@ -305,6 +305,20 @@
|
|
|
305
305
|
responses: [
|
|
306
306
|
{ code: '200 OK', description: 'Successful response.', body: "{...}" }
|
|
307
307
|
]
|
|
308
|
+
},
|
|
309
|
+
{
|
|
310
|
+
id: 'post-reload-worker-configs',
|
|
311
|
+
name: 'Reload Worker Configurations',
|
|
312
|
+
method: 'POST',
|
|
313
|
+
path: '/api/{version}/admin/reload-workers',
|
|
314
|
+
description: 'Triggers a dynamic reload of worker configurations from the TOML file. Requires client authentication.',
|
|
315
|
+
parameters: [
|
|
316
|
+
{ name: 'version', type: 'string', description: 'API Version', example: 'v1' }
|
|
317
|
+
],
|
|
318
|
+
request: { body: null },
|
|
319
|
+
responses: [
|
|
320
|
+
{ code: '200 OK', description: 'Successful response.', body: { "status": "worker_configs_reloaded" } }
|
|
321
|
+
]
|
|
308
322
|
}
|
|
309
323
|
]
|
|
310
324
|
},
|
|
@@ -178,10 +178,13 @@ class StateMachineBlueprint:
|
|
|
178
178
|
def render_graph(self, output_filename: Optional[str] = None, output_format: str = "png"):
|
|
179
179
|
import ast
|
|
180
180
|
import inspect
|
|
181
|
+
import logging
|
|
181
182
|
import textwrap
|
|
182
183
|
|
|
183
184
|
from graphviz import Digraph # type: ignore[import]
|
|
184
185
|
|
|
186
|
+
logger = logging.getLogger(__name__)
|
|
187
|
+
|
|
185
188
|
dot = Digraph(comment=f"State Machine for {self.name}")
|
|
186
189
|
dot.attr("node", shape="box", style="rounded")
|
|
187
190
|
all_handlers = list(self.handlers.items()) + [(ch.state, ch.func) for ch in self.conditional_handlers]
|
|
@@ -222,7 +225,11 @@ class StateMachineBlueprint:
|
|
|
222
225
|
value,
|
|
223
226
|
label=f"on {key}",
|
|
224
227
|
)
|
|
225
|
-
except (TypeError, OSError):
|
|
228
|
+
except (TypeError, OSError) as e:
|
|
229
|
+
logger.warning(
|
|
230
|
+
f"Could not parse handler '{handler_func.__name__}' for state '{handler_state}'. "
|
|
231
|
+
f"Graph may be incomplete. Error: {e}"
|
|
232
|
+
)
|
|
226
233
|
pass
|
|
227
234
|
for state in states:
|
|
228
235
|
dot.node(state, state)
|
|
@@ -584,6 +584,18 @@ class OrchestratorEngine:
|
|
|
584
584
|
jobs = await self.storage.get_quarantined_jobs()
|
|
585
585
|
return web.json_response(jobs)
|
|
586
586
|
|
|
587
|
+
async def _reload_worker_configs_handler(self, request: web.Request) -> web.Response:
|
|
588
|
+
"""Handles the dynamic reloading of worker configurations."""
|
|
589
|
+
logger.info("Received request to reload worker configurations.")
|
|
590
|
+
if not self.config.WORKERS_CONFIG_PATH:
|
|
591
|
+
return web.json_response(
|
|
592
|
+
{"error": "WORKERS_CONFIG_PATH is not set, cannot reload configs."},
|
|
593
|
+
status=400,
|
|
594
|
+
)
|
|
595
|
+
|
|
596
|
+
await load_worker_configs_to_redis(self.storage, self.config.WORKERS_CONFIG_PATH)
|
|
597
|
+
return web.json_response({"status": "worker_configs_reloaded"})
|
|
598
|
+
|
|
587
599
|
async def _flush_db_handler(self, request: web.Request) -> web.Response:
|
|
588
600
|
logger.warning("Received request to flush the database.")
|
|
589
601
|
await self.storage.flush_all()
|
|
@@ -643,6 +655,7 @@ class OrchestratorEngine:
|
|
|
643
655
|
app.router.add_get("/workers", self._get_workers_handler)
|
|
644
656
|
app.router.add_get("/jobs", self._get_jobs_handler)
|
|
645
657
|
app.router.add_get("/dashboard", self._get_dashboard_handler)
|
|
658
|
+
app.router.add_post("/admin/reload-workers", self._reload_worker_configs_handler)
|
|
646
659
|
|
|
647
660
|
if has_unversioned_routes:
|
|
648
661
|
self.app.add_subapp("/api/", protected_app)
|
|
@@ -1,3 +1,4 @@
|
|
|
1
|
+
from hashlib import sha256
|
|
1
2
|
from typing import Any, Awaitable, Callable
|
|
2
3
|
|
|
3
4
|
from aiohttp import web
|
|
@@ -89,9 +90,10 @@ def worker_auth_middleware_factory(
|
|
|
89
90
|
)
|
|
90
91
|
|
|
91
92
|
# --- Individual Token Check ---
|
|
92
|
-
|
|
93
|
-
if
|
|
94
|
-
|
|
93
|
+
expected_token_hash = await storage.get_worker_token(worker_id)
|
|
94
|
+
if expected_token_hash:
|
|
95
|
+
hashed_provided_token = sha256(provided_token.encode()).hexdigest()
|
|
96
|
+
if hashed_provided_token == expected_token_hash:
|
|
95
97
|
request["worker_id"] = worker_id # Attach authenticated worker_id
|
|
96
98
|
return await handler(request)
|
|
97
99
|
else:
|
|
@@ -1,3 +1,4 @@
|
|
|
1
|
+
from hashlib import sha256
|
|
1
2
|
from logging import getLogger
|
|
2
3
|
from os.path import exists
|
|
3
4
|
from tomllib import load
|
|
@@ -36,8 +37,10 @@ async def load_worker_configs_to_redis(storage: StorageBackend, config_path: str
|
|
|
36
37
|
continue
|
|
37
38
|
|
|
38
39
|
try:
|
|
40
|
+
# Hash the token before storing it
|
|
41
|
+
hashed_token = sha256(token.encode()).hexdigest()
|
|
39
42
|
# Store the token in a way that's easily retrievable by worker_id
|
|
40
|
-
await storage.set_worker_token(worker_id,
|
|
43
|
+
await storage.set_worker_token(worker_id, hashed_token)
|
|
41
44
|
logger.info(f"Loaded token for worker_id '{worker_id}'.")
|
|
42
45
|
except Exception as e:
|
|
43
46
|
logger.error(f"Failed to store token for worker_id '{worker_id}' in Redis: {e}")
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: avtomatika
|
|
3
|
-
Version: 1.
|
|
3
|
+
Version: 1.0b2
|
|
4
4
|
Summary: A state-machine based orchestrator for long-running jobs.
|
|
5
5
|
Project-URL: Homepage, https://github.com/avtomatika-ai/avtomatika
|
|
6
6
|
Project-URL: Bug Tracker, https://github.com/avtomatika-ai/avtomatika/issues
|
|
@@ -251,6 +251,11 @@ async def handle_normal(actions):
|
|
|
251
251
|
actions.transition_to("normal_processing")
|
|
252
252
|
```
|
|
253
253
|
|
|
254
|
+
> **Note on Limitations:** The current version of `.when()` uses a simple parser with the following limitations:
|
|
255
|
+
> * **No Nested Attributes:** You can only access direct fields of `context.initial_data` or `context.state_history` (e.g., `context.initial_data.field`). Nested objects (e.g., `context.initial_data.area.field`) are not supported.
|
|
256
|
+
> * **Simple Comparisons Only:** Only the following operators are supported: `==`, `!=`, `>`, `<`, `>=`, `<=`. Complex logical expressions with `AND`, `OR`, or `NOT` are not allowed.
|
|
257
|
+
> * **Limited Value Types:** The parser only recognizes strings (in quotes), integers, and floats. Boolean values (`True`, `False`) and `None` are not correctly parsed and will be treated as strings.
|
|
258
|
+
|
|
254
259
|
### 2. Delegating Tasks to Workers (`dispatch_task`)
|
|
255
260
|
|
|
256
261
|
This is the primary function for delegating work. The orchestrator will queue the task and wait for a worker to pick it up and return a result.
|
|
@@ -368,7 +373,9 @@ The orchestrator uses tokens to authenticate API requests.
|
|
|
368
373
|
* **Client Authentication**: All API clients must provide a token in the `X-Avtomatika-Token` header. The orchestrator validates this token against client configurations.
|
|
369
374
|
* **Worker Authentication**: Workers must provide a token in the `X-Worker-Token` header.
|
|
370
375
|
* `GLOBAL_WORKER_TOKEN`: You can set a global token for all workers using this environment variable. For development and testing, it defaults to `"secure-worker-token"`.
|
|
371
|
-
* **Individual Tokens**: For production, it is recommended to define individual tokens for each worker in a separate configuration file and provide its path via the `WORKERS_CONFIG_PATH` environment variable.
|
|
376
|
+
* **Individual Tokens**: For production, it is recommended to define individual tokens for each worker in a separate configuration file and provide its path via the `WORKERS_CONFIG_PATH` environment variable. Tokens from this file are stored in a hashed format for security.
|
|
377
|
+
|
|
378
|
+
> **Note on Dynamic Reloading:** The worker configuration file can be reloaded without restarting the orchestrator by sending an authenticated `POST` request to the `/api/v1/admin/reload-workers` endpoint. This allows for dynamic updates of worker tokens.
|
|
372
379
|
|
|
373
380
|
### Observability
|
|
374
381
|
|
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
from unittest.mock import Mock
|
|
2
|
+
|
|
3
|
+
import pytest
|
|
4
|
+
import zstandard
|
|
5
|
+
from aiohttp import web
|
|
6
|
+
from src.avtomatika.compression import _compress_gzip, compression_middleware
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
@pytest.mark.asyncio
|
|
10
|
+
async def test_no_compression_for_small_body():
|
|
11
|
+
"""Ensures responses with a body smaller than 500 bytes are not compressed."""
|
|
12
|
+
request = Mock()
|
|
13
|
+
request.headers = {"Accept-Encoding": "zstd"}
|
|
14
|
+
|
|
15
|
+
# The handler returns a response with a small body
|
|
16
|
+
async def handler(req):
|
|
17
|
+
return web.Response(body=b"small body")
|
|
18
|
+
|
|
19
|
+
response = await compression_middleware(request, handler)
|
|
20
|
+
assert "Content-Encoding" not in response.headers
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@pytest.mark.asyncio
|
|
24
|
+
async def test_no_compression_if_already_encoded():
|
|
25
|
+
"""Ensures responses that already have a Content-Encoding are not compressed again."""
|
|
26
|
+
request = Mock()
|
|
27
|
+
request.headers = {"Accept-Encoding": "zstd"}
|
|
28
|
+
|
|
29
|
+
async def handler(req):
|
|
30
|
+
response = web.Response(body=b"i am already compressed" * 100)
|
|
31
|
+
response.headers["Content-Encoding"] = "br"
|
|
32
|
+
return response
|
|
33
|
+
|
|
34
|
+
response = await compression_middleware(request, handler)
|
|
35
|
+
assert response.headers["Content-Encoding"] == "br"
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
@pytest.mark.asyncio
|
|
39
|
+
async def test_websocket_response_is_ignored():
|
|
40
|
+
"""Ensures WebSocket responses are not compressed."""
|
|
41
|
+
request = Mock()
|
|
42
|
+
request.headers = {"Accept-Encoding": "zstd"}
|
|
43
|
+
|
|
44
|
+
async def handler(req):
|
|
45
|
+
ws_response = web.WebSocketResponse()
|
|
46
|
+
# In a real scenario, prepare() would be called, but for middleware testing,
|
|
47
|
+
# returning the instance is sufficient.
|
|
48
|
+
return ws_response
|
|
49
|
+
|
|
50
|
+
response = await compression_middleware(request, handler)
|
|
51
|
+
assert isinstance(response, web.WebSocketResponse)
|
|
52
|
+
assert "Content-Encoding" not in response.headers
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
@pytest.mark.asyncio
|
|
56
|
+
async def test_zstd_compression_occurs():
|
|
57
|
+
"""Tests that zstd compression is applied correctly."""
|
|
58
|
+
request = Mock()
|
|
59
|
+
request.headers = {"Accept-Encoding": "zstd"}
|
|
60
|
+
large_body = b"some large body content" * 100
|
|
61
|
+
|
|
62
|
+
async def handler(req):
|
|
63
|
+
return web.Response(body=large_body)
|
|
64
|
+
|
|
65
|
+
response = await compression_middleware(request, handler)
|
|
66
|
+
assert response.headers["Content-Encoding"] == "zstd"
|
|
67
|
+
|
|
68
|
+
decompressor = zstandard.ZstdDecompressor()
|
|
69
|
+
decompressed_body = decompressor.decompress(response.body)
|
|
70
|
+
assert decompressed_body == large_body
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
@pytest.mark.asyncio
|
|
74
|
+
async def test_gzip_compression_occurs():
|
|
75
|
+
"""Tests that gzip compression is applied correctly."""
|
|
76
|
+
request = Mock()
|
|
77
|
+
request.headers = {"Accept-Encoding": "gzip"}
|
|
78
|
+
large_body = b"some large body content for gzip" * 100
|
|
79
|
+
|
|
80
|
+
async def handler(req):
|
|
81
|
+
return web.Response(body=large_body)
|
|
82
|
+
|
|
83
|
+
response = await compression_middleware(request, handler)
|
|
84
|
+
assert response.headers["Content-Encoding"] == "gzip"
|
|
85
|
+
|
|
86
|
+
import gzip
|
|
87
|
+
|
|
88
|
+
decompressed_body = gzip.decompress(response.body)
|
|
89
|
+
assert decompressed_body == large_body
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
@pytest.mark.asyncio
|
|
93
|
+
async def test_compression_failure_returns_original_response():
|
|
94
|
+
"""Tests that if the compression function fails, the original response is returned."""
|
|
95
|
+
request = Mock()
|
|
96
|
+
request.headers = {"Accept-Encoding": "gzip"}
|
|
97
|
+
large_body = b"some large body that will fail to compress" * 100
|
|
98
|
+
|
|
99
|
+
# Mock the compression function to raise an exception
|
|
100
|
+
original_compress = _compress_gzip
|
|
101
|
+
try:
|
|
102
|
+
|
|
103
|
+
def failing_compress_gzip(data):
|
|
104
|
+
raise ValueError("Compression failed!")
|
|
105
|
+
|
|
106
|
+
# Monkeypatch the function
|
|
107
|
+
import src.avtomatika.compression
|
|
108
|
+
|
|
109
|
+
src.avtomatika.compression._compress_gzip = failing_compress_gzip
|
|
110
|
+
|
|
111
|
+
async def handler(req):
|
|
112
|
+
return web.Response(body=large_body)
|
|
113
|
+
|
|
114
|
+
response = await compression_middleware(request, handler)
|
|
115
|
+
assert "Content-Encoding" not in response.headers
|
|
116
|
+
assert response.body == large_body
|
|
117
|
+
finally:
|
|
118
|
+
# Restore the original function
|
|
119
|
+
import src.avtomatika.compression
|
|
120
|
+
|
|
121
|
+
src.avtomatika.compression._compress_gzip = original_compress
|
|
@@ -1,4 +1,5 @@
|
|
|
1
1
|
import asyncio
|
|
2
|
+
import hashlib
|
|
2
3
|
import json
|
|
3
4
|
import os
|
|
4
5
|
from unittest.mock import AsyncMock
|
|
@@ -491,7 +492,12 @@ async def test_task_cancellation_via_websocket_mocked(aiohttp_client, app):
|
|
|
491
492
|
|
|
492
493
|
@pytest.mark.parametrize(
|
|
493
494
|
"app",
|
|
494
|
-
[
|
|
495
|
+
[
|
|
496
|
+
{
|
|
497
|
+
"extra_blueprints": [cancellation_bp],
|
|
498
|
+
"workers_config_path": os.path.join(os.path.dirname(os.path.abspath(__file__)), "workers.toml"),
|
|
499
|
+
}
|
|
500
|
+
],
|
|
495
501
|
indirect=True,
|
|
496
502
|
)
|
|
497
503
|
@pytest.mark.asyncio
|
|
@@ -502,7 +508,8 @@ async def test_worker_individual_token_auth(aiohttp_client, app):
|
|
|
502
508
|
|
|
503
509
|
worker_id = "worker-with-individual-token"
|
|
504
510
|
individual_token = "individual-secret-for-worker-1"
|
|
505
|
-
|
|
511
|
+
hashed_individual_token = hashlib.sha256(individual_token.encode()).hexdigest()
|
|
512
|
+
await storage.set_worker_token(worker_id, hashed_individual_token)
|
|
506
513
|
|
|
507
514
|
headers = {"X-Worker-Token": individual_token}
|
|
508
515
|
payload = {"worker_id": worker_id, "worker_type": "test", "supported_tasks": ["test"]}
|
|
@@ -1,3 +1,4 @@
|
|
|
1
|
+
import hashlib
|
|
1
2
|
import os
|
|
2
3
|
from unittest.mock import AsyncMock
|
|
3
4
|
|
|
@@ -23,9 +24,12 @@ token = "token-2"
|
|
|
23
24
|
|
|
24
25
|
await load_worker_configs_to_redis(storage, config_path)
|
|
25
26
|
|
|
27
|
+
hashed_token_1 = hashlib.sha256(b"token-1").hexdigest()
|
|
28
|
+
hashed_token_2 = hashlib.sha256(b"token-2").hexdigest()
|
|
29
|
+
|
|
26
30
|
assert storage.set_worker_token.call_count == 2
|
|
27
|
-
storage.set_worker_token.assert_any_call("worker-1",
|
|
28
|
-
storage.set_worker_token.assert_any_call("worker-2",
|
|
31
|
+
storage.set_worker_token.assert_any_call("worker-1", hashed_token_1)
|
|
32
|
+
storage.set_worker_token.assert_any_call("worker-2", hashed_token_2)
|
|
29
33
|
|
|
30
34
|
os.remove(config_path)
|
|
31
35
|
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|