Skip to content

Commit a9736d6

Browse files
committed
feat: add operational utilities
- Add tools/testing/publish_amqp.py for AMQP message publishing - Add validate-setup.sh for environment validation - CLI tools for testing and troubleshooting - Add test suite
1 parent da5ac25 commit a9736d6

File tree

3 files changed

+406
-0
lines changed

3 files changed

+406
-0
lines changed

tools/testing/publish_amqp.py

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
#!/usr/bin/env python3
2+
"""AMQP message publisher for triggering GeoZarr conversion workflows.
3+
4+
Publishes JSON payloads to RabbitMQ exchanges with support for
5+
dynamic routing key templates based on payload fields.
6+
"""
7+
8+
from __future__ import annotations
9+
10+
import argparse
11+
import json
12+
import logging
13+
import sys
14+
from pathlib import Path
15+
from typing import Any
16+
17+
import pika
18+
from tenacity import retry, stop_after_attempt, wait_exponential
19+
20+
logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
21+
logger = logging.getLogger(__name__)
22+
23+
24+
def load_payload(payload_file: Path) -> dict[str, Any]:
25+
"""Load JSON payload from file."""
26+
try:
27+
data: dict[str, Any] = json.loads(payload_file.read_text())
28+
return data
29+
except FileNotFoundError:
30+
logger.exception("Payload file not found", extra={"file": str(payload_file)})
31+
sys.exit(1)
32+
except json.JSONDecodeError:
33+
logger.exception("Invalid JSON in payload file", extra={"file": str(payload_file)})
34+
sys.exit(1)
35+
36+
37+
def format_routing_key(template: str, payload: dict[str, Any]) -> str:
38+
"""Format routing key template using payload fields.
39+
40+
Example: "eopf.item.found.{collection}" → "eopf.item.found.sentinel-2-l2a"
41+
"""
42+
try:
43+
return template.format(**payload)
44+
except KeyError:
45+
logger.exception(
46+
"Missing required field in payload for routing key template",
47+
extra={"template": template, "available_fields": list(payload.keys())},
48+
)
49+
sys.exit(1)
50+
51+
52+
@retry(stop=stop_after_attempt(3), wait=wait_exponential(min=1, max=10))
53+
def publish_message(
54+
host: str,
55+
port: int,
56+
user: str,
57+
password: str,
58+
exchange: str,
59+
routing_key: str,
60+
payload: dict[str, Any],
61+
virtual_host: str = "/",
62+
) -> None:
63+
"""Publish message to RabbitMQ exchange with automatic retry."""
64+
credentials = pika.PlainCredentials(user, password)
65+
parameters = pika.ConnectionParameters(
66+
host=host,
67+
port=port,
68+
virtual_host=virtual_host,
69+
credentials=credentials,
70+
)
71+
72+
logger.info("Connecting to amqp://%s@%s:%s%s", user, host, port, virtual_host)
73+
connection = pika.BlockingConnection(parameters)
74+
try:
75+
channel = connection.channel()
76+
channel.basic_publish(
77+
exchange=exchange,
78+
routing_key=routing_key,
79+
body=json.dumps(payload),
80+
properties=pika.BasicProperties(
81+
content_type="application/json",
82+
delivery_mode=2,
83+
),
84+
)
85+
logger.info("Published to exchange='%s' routing_key='%s'", exchange, routing_key)
86+
logger.debug("Payload: %s", json.dumps(payload, indent=2))
87+
finally:
88+
connection.close()
89+
90+
91+
def main() -> None:
92+
"""CLI entry point for AMQP message publisher."""
93+
parser = argparse.ArgumentParser(
94+
description="Publish JSON payload to RabbitMQ exchange for workflow triggers"
95+
)
96+
parser.add_argument("--host", required=True, help="RabbitMQ host")
97+
parser.add_argument("--port", type=int, default=5672, help="RabbitMQ port")
98+
parser.add_argument("--user", required=True, help="RabbitMQ username")
99+
parser.add_argument("--password", required=True, help="RabbitMQ password")
100+
parser.add_argument("--virtual-host", default="/", help="RabbitMQ virtual host")
101+
parser.add_argument("--exchange", required=True, help="RabbitMQ exchange name")
102+
parser.add_argument("--routing-key", help="Static routing key")
103+
parser.add_argument(
104+
"--routing-key-template",
105+
help="Template with {field} placeholders (e.g., 'eopf.item.found.{collection}')",
106+
)
107+
parser.add_argument("--payload-file", type=Path, required=True, help="JSON payload file path")
108+
109+
args = parser.parse_args()
110+
111+
if not args.routing_key and not args.routing_key_template:
112+
parser.error("Must provide either --routing-key or --routing-key-template")
113+
if args.routing_key and args.routing_key_template:
114+
parser.error("Cannot use both --routing-key and --routing-key-template")
115+
116+
payload = load_payload(args.payload_file)
117+
routing_key = args.routing_key or format_routing_key(args.routing_key_template, payload)
118+
119+
try:
120+
publish_message(
121+
host=args.host,
122+
port=args.port,
123+
user=args.user,
124+
password=args.password,
125+
exchange=args.exchange,
126+
routing_key=routing_key,
127+
payload=payload,
128+
virtual_host=args.virtual_host,
129+
)
130+
except Exception:
131+
logger.exception(
132+
"Failed to publish AMQP message",
133+
extra={
134+
"exchange": args.exchange,
135+
"routing_key": routing_key,
136+
"host": args.host,
137+
},
138+
)
139+
sys.exit(1)
140+
141+
142+
if __name__ == "__main__":
143+
main()

tools/testing/test_publish_amqp.py

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
"""Unit tests for publish_amqp.py script."""
2+
3+
from __future__ import annotations
4+
5+
import json
6+
import sys
7+
from pathlib import Path
8+
9+
import pika.exceptions
10+
import pytest
11+
12+
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "scripts"))
13+
from publish_amqp import format_routing_key, load_payload
14+
15+
16+
@pytest.fixture
17+
def sample_payload() -> dict[str, str]:
18+
"""Sample payload for tests."""
19+
return {"collection": "sentinel-2-l2a", "item_id": "test-123"}
20+
21+
22+
@pytest.fixture
23+
def payload_file(tmp_path: Path, sample_payload: dict[str, str]) -> Path:
24+
"""Create a temporary payload file."""
25+
file = tmp_path / "payload.json"
26+
file.write_text(json.dumps(sample_payload))
27+
return file
28+
29+
30+
class TestLoadPayload:
31+
"""Tests for payload loading."""
32+
33+
def test_valid_payload(self, payload_file: Path, sample_payload: dict[str, str]) -> None:
34+
"""Load valid JSON payload."""
35+
assert load_payload(payload_file) == sample_payload
36+
37+
def test_missing_file(self, tmp_path: Path) -> None:
38+
"""Handle missing file with exit code 1."""
39+
with pytest.raises(SystemExit, match="1"):
40+
load_payload(tmp_path / "missing.json")
41+
42+
def test_invalid_json(self, tmp_path: Path) -> None:
43+
"""Handle invalid JSON with exit code 1."""
44+
invalid = tmp_path / "invalid.json"
45+
invalid.write_text("{not valid json")
46+
with pytest.raises(SystemExit, match="1"):
47+
load_payload(invalid)
48+
49+
50+
class TestFormatRoutingKey:
51+
"""Tests for routing key formatting."""
52+
53+
@pytest.mark.parametrize(
54+
("template", "payload", "expected"),
55+
[
56+
(
57+
"eopf.item.found.{collection}",
58+
{"collection": "sentinel-2-l2a"},
59+
"eopf.item.found.sentinel-2-l2a",
60+
),
61+
(
62+
"{env}.{service}.{collection}",
63+
{"env": "prod", "service": "ingest", "collection": "s1"},
64+
"prod.ingest.s1",
65+
),
66+
("static.key", {"collection": "sentinel-2"}, "static.key"),
67+
],
68+
)
69+
def test_format_templates(self, template: str, payload: dict[str, str], expected: str) -> None:
70+
"""Format various routing key templates."""
71+
assert format_routing_key(template, payload) == expected
72+
73+
def test_missing_field(self) -> None:
74+
"""Handle missing field with exit code 1."""
75+
with pytest.raises(SystemExit, match="1"):
76+
format_routing_key("eopf.item.found.{collection}", {"item_id": "test"})
77+
78+
79+
class TestPublishMessage:
80+
"""Tests for message publishing (mocked)."""
81+
82+
def test_publish_success(self, mocker: pytest.MonkeyPatch) -> None:
83+
"""Publish message successfully."""
84+
from publish_amqp import publish_message
85+
86+
mock_conn = mocker.patch("publish_amqp.pika.BlockingConnection")
87+
mock_channel = mocker.MagicMock()
88+
mock_conn.return_value.channel.return_value = mock_channel
89+
90+
publish_message(
91+
host="rabbitmq.test",
92+
port=5672,
93+
user="testuser",
94+
password="testpass",
95+
exchange="test_exchange",
96+
routing_key="test.key",
97+
payload={"test": "data"},
98+
)
99+
100+
mock_conn.assert_called_once()
101+
mock_channel.basic_publish.assert_called_once()
102+
call = mock_channel.basic_publish.call_args.kwargs
103+
assert call["exchange"] == "test_exchange"
104+
assert call["routing_key"] == "test.key"
105+
assert json.loads(call["body"]) == {"test": "data"}
106+
107+
def test_connection_retry(self, mocker: pytest.MonkeyPatch) -> None:
108+
"""Verify tenacity retry on transient failures."""
109+
from publish_amqp import publish_message
110+
111+
mock_conn = mocker.patch("publish_amqp.pika.BlockingConnection")
112+
mock_channel = mocker.MagicMock()
113+
114+
# Fail twice, succeed on third attempt
115+
mock_conn.side_effect = [
116+
pika.exceptions.AMQPConnectionError("Transient error"),
117+
pika.exceptions.AMQPConnectionError("Transient error"),
118+
mocker.MagicMock(channel=mocker.MagicMock(return_value=mock_channel)),
119+
]
120+
121+
publish_message(
122+
host="rabbitmq.test",
123+
port=5672,
124+
user="testuser",
125+
password="testpass",
126+
exchange="test_exchange",
127+
routing_key="test.key",
128+
payload={"test": "data"},
129+
)
130+
131+
assert mock_conn.call_count == 3

0 commit comments

Comments
 (0)