Add unit tests to the MCP server (#12)

* added tests

* all tests passed

Author: kenliao94

src/tests/__init__.py

@@ -0,0 +1,13 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

src/tests/conftest.py

@@ -0,0 +1,36 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+
+@pytest.fixture
+def sample_queue_data():
+    return [
+        {"name": "queue1", "messages": 10, "vhost": "/"},
+        {"name": "queue2", "messages": 5, "vhost": "/"},
+    ]
+
+
+@pytest.fixture
+def sample_exchange_data():
+    return [
+        {"name": "exchange1", "type": "direct", "vhost": "/"},
+        {"name": "exchange2", "type": "fanout", "vhost": "/"},
+    ]
+
+
+@pytest.fixture
+def sample_vhost_data():
+    return [{"name": "/"}, {"name": "test-vhost"}]

src/tests/test_admin.py

@@ -0,0 +1,105 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from src.rabbitmq.admin import RabbitMQAdmin
+
+
+@pytest.fixture
+def admin():
+    return RabbitMQAdmin("localhost", "user", "pass")
+
+
+@pytest.fixture
+def mock_response():
+    response = MagicMock()
+    response.status_code = 200
+    response.json.return_value = {"test": "data"}
+    return response
+
+
+class TestRabbitMQAdmin:
+    def test_init(self, admin):
+        assert admin.protocol == "https"
+        assert admin.base_url == "https://localhost/api"
+
+    @patch("src.rabbitmq.admin.requests.request")
+    def test_list_queues(self, mock_request, admin, mock_response):
+        mock_request.return_value = mock_response
+        result = admin.list_queues()
+        assert result == {"test": "data"}
+        mock_request.assert_called_once()
+
+    @patch("src.rabbitmq.admin.requests.request")
+    def test_get_queue_info(self, mock_request, admin, mock_response):
+        mock_request.return_value = mock_response
+        result = admin.get_queue_info("test-queue")
+        assert result == {"test": "data"}
+
+    @patch("src.rabbitmq.admin.requests.request")
+    def test_delete_queue(self, mock_request, admin, mock_response):
+        mock_request.return_value = mock_response
+        admin.delete_queue("test-queue")
+        mock_request.assert_called_once()
+
+    @patch("src.rabbitmq.admin.requests.request")
+    def test_purge_queue(self, mock_request, admin, mock_response):
+        mock_request.return_value = mock_response
+        admin.purge_queue("test-queue")
+        mock_request.assert_called_once()
+
+    @patch("src.rabbitmq.admin.requests.request")
+    def test_list_exchanges(self, mock_request, admin, mock_response):
+        mock_request.return_value = mock_response
+        result = admin.list_exchanges()
+        assert result == {"test": "data"}
+
+    @patch("src.rabbitmq.admin.requests.request")
+    def test_get_exchange_info(self, mock_request, admin, mock_response):
+        mock_request.return_value = mock_response
+        result = admin.get_exchange_info("test-exchange")
+        assert result == {"test": "data"}
+
+    @patch("src.rabbitmq.admin.requests.request")
+    def test_delete_exchange(self, mock_request, admin, mock_response):
+        mock_request.return_value = mock_response
+        admin.delete_exchange("test-exchange")
+        mock_request.assert_called_once()
+
+    @patch("src.rabbitmq.admin.requests.request")
+    def test_get_overview(self, mock_request, admin, mock_response):
+        mock_request.return_value = mock_response
+        result = admin.get_overview()
+        assert result == {"test": "data"}
+
+    @patch("src.rabbitmq.admin.requests.request")
+    def test_list_vhosts(self, mock_request, admin, mock_response):
+        mock_request.return_value = mock_response
+        result = admin.list_vhosts()
+        assert result == {"test": "data"}
+
+    @patch("src.rabbitmq.admin.requests.request")
+    def test_get_alarm_status(self, mock_request, admin, mock_response):
+        mock_request.return_value = mock_response
+        result = admin.get_alarm_status()
+        assert result == 200
+
+    @patch("src.rabbitmq.admin.requests.request")
+    def test_get_broker_definition(self, mock_request, admin, mock_response):
+        mock_request.return_value = mock_response
+        result = admin.get_broker_definition()
+        assert result == {"test": "data"}

src/tests/test_connection.py

@@ -0,0 +1,54 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from src.rabbitmq.connection import RabbitMQConnection, validate_rabbitmq_name
+
+
+class TestRabbitMQConnection:
+    def test_init_with_tls(self):
+        conn = RabbitMQConnection("localhost", "user", "pass", use_tls=True)
+        assert conn.protocol == "amqps"
+        assert "amqps://user:pass@localhost:5671" == conn.url
+
+    def test_init_without_tls(self):
+        conn = RabbitMQConnection("localhost", "user", "pass", use_tls=False)
+        assert conn.protocol == "amqp"
+        assert "amqp://user:pass@localhost:5671" == conn.url
+
+
+class TestValidateRabbitMQName:
+    def test_valid_names(self):
+        validate_rabbitmq_name("valid-queue", "Queue")
+        validate_rabbitmq_name("valid_queue", "Queue")
+        validate_rabbitmq_name("valid.queue", "Queue")
+        validate_rabbitmq_name("valid:queue", "Queue")
+        validate_rabbitmq_name("queue123", "Queue")
+
+    def test_empty_name(self):
+        with pytest.raises(ValueError, match="cannot be empty"):
+            validate_rabbitmq_name("", "Queue")
+
+    def test_whitespace_only(self):
+        with pytest.raises(ValueError, match="cannot be empty"):
+            validate_rabbitmq_name("   ", "Queue")
+
+    def test_invalid_characters(self):
+        with pytest.raises(ValueError, match="can only contain"):
+            validate_rabbitmq_name("invalid@queue", "Queue")
+
+    def test_too_long(self):
+        with pytest.raises(ValueError, match="must be less than 255"):
+            validate_rabbitmq_name("a" * 256, "Queue")

src/tests/test_handlers.py

@@ -0,0 +1,173 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from src.rabbitmq.handlers import (
+    handle_delete_exchange,
+    handle_delete_queue,
+    handle_get_cluster_nodes,
+    handle_get_definition,
+    handle_get_exchange_info,
+    handle_get_guidelines,
+    handle_get_queue_info,
+    handle_is_broker_in_alarm,
+    handle_is_node_in_quorum_critical,
+    handle_list_connections,
+    handle_list_consumers,
+    handle_list_exchanges,
+    handle_list_queues,
+    handle_list_shovels,
+    handle_list_users,
+    handle_list_vhosts,
+    handle_purge_queue,
+    handle_shovel,
+)
+
+
+@pytest.fixture
+def mock_admin():
+    admin = MagicMock()
+    admin.list_queues.return_value = [{"name": "queue1"}, {"name": "queue2"}]
+    admin.list_exchanges.return_value = [{"name": "exchange1"}]
+    admin.list_vhosts.return_value = [{"name": "vhost1"}]
+    admin.get_queue_info.return_value = {"name": "queue1", "messages": 10}
+    admin.get_exchange_info.return_value = {"name": "exchange1", "type": "direct"}
+    admin.get_alarm_status.return_value = 200
+    admin.get_is_node_quorum_critical.return_value = 200
+    admin.get_broker_definition.return_value = {"queues": []}
+    admin.list_shovels.return_value = [{"name": "shovel1"}]
+    admin.get_shovel_info.return_value = {"name": "shovel1"}
+    admin.list_connections.return_value = [
+        {
+            "auth_mechanism": "PLAIN",
+            "channels": 1,
+            "client_properties": {},
+            "connected_at": 1000000,
+            "state": "running",
+        }
+    ]
+    admin.list_consumers.return_value = [{"queue": "queue1"}]
+    admin.list_users.return_value = [{"name": "user1"}]
+    admin.get_cluster_nodes.return_value = [
+        {
+            "name": "node1",
+            "mem_alarm": False,
+            "disk_free_alarm": False,
+            "disk_free": 1000,
+            "mem_limit": 2000,
+            "mem_used": 1000,
+            "rates_mode": "basic",
+            "uptime": 10000,
+            "running": True,
+            "queue_created": 5,
+            "queue_deleted": 1,
+            "connection_created": 10,
+        }
+    ]
+    return admin
+
+
+class TestHandlers:
+    def test_handle_list_queues(self, mock_admin):
+        result = handle_list_queues(mock_admin)
+        assert result == ["queue1", "queue2"]
+
+    def test_handle_list_exchanges(self, mock_admin):
+        result = handle_list_exchanges(mock_admin)
+        assert result == ["exchange1"]
+
+    def test_handle_list_vhosts(self, mock_admin):
+        result = handle_list_vhosts(mock_admin)
+        assert result == ["vhost1"]
+
+    def test_handle_get_queue_info(self, mock_admin):
+        result = handle_get_queue_info(mock_admin, "queue1")
+        assert result == {"name": "queue1", "messages": 10}
+
+    def test_handle_get_exchange_info(self, mock_admin):
+        result = handle_get_exchange_info(mock_admin, "exchange1")
+        assert result == {"name": "exchange1", "type": "direct"}
+
+    def test_handle_delete_queue(self, mock_admin):
+        handle_delete_queue(mock_admin, "queue1")
+        mock_admin.delete_queue.assert_called_once_with("queue1", "/")
+
+    def test_handle_purge_queue(self, mock_admin):
+        handle_purge_queue(mock_admin, "queue1")
+        mock_admin.purge_queue.assert_called_once_with("queue1", "/")
+
+    def test_handle_delete_exchange(self, mock_admin):
+        handle_delete_exchange(mock_admin, "exchange1")
+        mock_admin.delete_exchange.assert_called_once_with("exchange1", "/")
+
+    def test_handle_is_broker_in_alarm_false(self, mock_admin):
+        result = handle_is_broker_in_alarm(mock_admin)
+        assert result is False
+
+    def test_handle_is_broker_in_alarm_true(self, mock_admin):
+        mock_admin.get_alarm_status.return_value = 500
+        result = handle_is_broker_in_alarm(mock_admin)
+        assert result is True
+
+    def test_handle_is_node_in_quorum_critical_false(self, mock_admin):
+        result = handle_is_node_in_quorum_critical(mock_admin)
+        assert result is False
+
+    def test_handle_is_node_in_quorum_critical_true(self, mock_admin):
+        mock_admin.get_is_node_quorum_critical.return_value = 500
+        result = handle_is_node_in_quorum_critical(mock_admin)
+        assert result is True
+
+    def test_handle_get_definition(self, mock_admin):
+        result = handle_get_definition(mock_admin)
+        assert result == {"queues": []}
+
+    def test_handle_list_shovels(self, mock_admin):
+        result = handle_list_shovels(mock_admin)
+        assert result == [{"name": "shovel1"}]
+
+    def test_handle_shovel(self, mock_admin):
+        result = handle_shovel(mock_admin, "shovel1")
+        assert result == {"name": "shovel1"}
+
+    def test_handle_list_connections(self, mock_admin):
+        result = handle_list_connections(mock_admin)
+        assert len(result) == 1
+        assert result[0]["auth_mechanism"] == "PLAIN"
+
+    def test_handle_list_consumers(self, mock_admin):
+        result = handle_list_consumers(mock_admin)
+        assert result == [{"queue": "queue1"}]
+
+    def test_handle_list_users(self, mock_admin):
+        result = handle_list_users(mock_admin)
+        assert result == [{"name": "user1"}]
+
+    def test_handle_get_cluster_nodes(self, mock_admin):
+        result = handle_get_cluster_nodes(mock_admin)
+        assert len(result) == 1
+        assert result[0]["name"] == "node1"
+        assert result[0]["mem_used_in_percentage"] == 50.0
+
+    def test_handle_get_guidelines_valid(self):
+        result = handle_get_guidelines("rabbimq_broker_sizing_guide")
+        assert isinstance(result, str)
+        assert len(result) > 0
+
+    def test_handle_get_guidelines_invalid(self):
+        with pytest.raises(ValueError, match="doesn't exist"):
+            handle_get_guidelines("invalid_guide")