Increase test coverage

Author: Michael Christensen

src/amazon-keyspaces-mcp-server/awslabs/amazon_keyspaces_mcp_server/client.py

@@ -456,8 +456,8 @@ def _build_service_characteristics(self) -> Dict[str, Any]:
 
     def close(self) -> None:
         """Close the session."""
-        if hasattr(self, 'session') and self.session:
-            if self.session.cluster:
-                self.session.cluster.shutdown()
-            self.session.shutdown()
+        if self._session:
+            if self._session.cluster:
+                self._session.cluster.shutdown()
+            self._session.shutdown()
             logger.info('Closed session')

src/amazon-keyspaces-mcp-server/awslabs/amazon_keyspaces_mcp_server/services.py

@@ -72,7 +72,7 @@ async def execute_read_only_query(self, keyspace_name: str, query: str) -> Dict[
                                 + query[table_name_start:]
                             )
 
-        return self.cassandra_client.execute_read_only_query(full_query)
+        return await self.cassandra_client.execute_read_only_query(full_query)
 
 
 class SchemaService:

src/amazon-keyspaces-mcp-server/tests/__init__.py

@@ -1,12 +1,15 @@
 # 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. A copy of the License is located at
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
 #
 #    http://www.apache.org/licenses/LICENSE-2.0
 #
-# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
-# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# or in the 'license' file accompanying this file.
+# This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied.
+# See the License for the specific language governing permissions
 # and limitations under the License.
 """
 Test package for keyspaces-mcp.

src/amazon-keyspaces-mcp-server/tests/test_client.py

@@ -1,12 +1,15 @@
 # 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. A copy of the License is located at
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
 #
 #    http://www.apache.org/licenses/LICENSE-2.0
 #
-# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
-# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# or in the 'license' file accompanying this file.
+# This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied.
+# See the License for the specific language governing permissions
 # and limitations under the License.
 """Tests for the awslabs.amazon-keyspaces-mcp-server package."""
 
@@ -19,31 +22,23 @@ class TestInit:
 
     def test_version(self):
         """Test that __version__ is defined and follows semantic versioning."""
-        # Import the module
+        # pylint: disable=import-outside-toplevel
         import awslabs.amazon_keyspaces_mcp_server
 
-        # Check that __version__ is defined
         assert hasattr(awslabs.amazon_keyspaces_mcp_server, '__version__')
-
-        # Check that __version__ is a string
         assert isinstance(awslabs.amazon_keyspaces_mcp_server.__version__, str)
 
-        # Check that __version__ follows semantic versioning (major.minor.patch)
         version_pattern = r'^\d+\.\d+\.\d+$'
-        assert re.match(version_pattern, awslabs.amazon_keyspaces_mcp_server.__version__), (
-            f"Version '{awslabs.amazon_keyspaces_mcp_server.__version__}' does not follow semantic versioning"
+        version = awslabs.amazon_keyspaces_mcp_server.__version__
+        assert re.match(version_pattern, version), (
+            f"Version '{version}' does not follow semantic versioning"
         )
 
     def test_module_reload(self):
         """Test that the module can be reloaded."""
-        # Import the module
+        # pylint: disable=import-outside-toplevel
         import awslabs.amazon_keyspaces_mcp_server
 
-        # Store the original version
         original_version = awslabs.amazon_keyspaces_mcp_server.__version__
-
-        # Reload the module
         importlib.reload(awslabs.amazon_keyspaces_mcp_server)
-
-        # Check that the version is still the same
         assert awslabs.amazon_keyspaces_mcp_server.__version__ == original_version

src/amazon-keyspaces-mcp-server/tests/test_init.py

@@ -1,18 +1,23 @@
 # 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. A copy of the License is located at
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+# A copy of the License is located at
 #
 #    http://www.apache.org/licenses/LICENSE-2.0
 #
-# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
-# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# or in the 'license' file accompanying this file.
+# This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied.
+# See the License for the specific language governing permissions
 # and limitations under the License.
 """Tests for the main function in server.py."""
 
-from awslabs.amazon_keyspaces_mcp_server.server import main
+import inspect
 from unittest.mock import patch
 
+from awslabs.amazon_keyspaces_mcp_server.server import main
+
 
 class TestMain:
     """Tests for the main function."""
@@ -27,26 +32,16 @@ def test_main_default(self, mock_run):
         2. The mcp.run method is called once
         3. No transport parameter is passed to mcp.run
         """
-        # Call the main function
         main()
-
-        # Check that mcp.run was called with the correct arguments
         mock_run.assert_called_once()
         assert mock_run.call_args[1].get('transport') is None
 
     def test_module_execution(self):
         """Test the module execution when run as __main__."""
-        # This test directly executes the code in the if __name__ == '__main__': block
-        # to ensure coverage of that line
-
-        # Get the source code of the module
-        import inspect
+        # pylint: disable=import-outside-toplevel
         from awslabs.amazon_keyspaces_mcp_server import server
 
-        # Get the source code
         source = inspect.getsource(server)
-
-        # Check that the module has the if __name__ == '__main__': block
         assert "if __name__ == '__main__':" in source
         assert 'main()' in source
 

src/amazon-keyspaces-mcp-server/tests/test_main.py

@@ -0,0 +1,35 @@
+# 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.
+
+"""Tests for models module."""
+
+import unittest
+
+from awslabs.amazon_keyspaces_mcp_server.models import QueryInput
+
+
+class TestModels(unittest.TestCase):
+    """Test cases for models."""
+
+    def test_sanitize_query_with_hidden_characters(self):
+        """Test query sanitization removes hidden unicode characters."""
+        input_data = QueryInput(
+            keyspace='test',
+            query='SELECT\u200B * FROM\uFEFF users\u0000'
+        )
+        self.assertEqual(input_data.query, 'SELECT * FROM users')
+
+
+if __name__ == '__main__':
+    unittest.main()