Add test script for schema registry and query/search agent tools

test_schema_tools.py

@@ -0,0 +1,88 @@
+"""Test script for schema registry and query/search agent tools
+
+Tests mapping, edge cases, error handling, and rollback/versioning with logs.
+"""
+import unittest
+from schema_registry import SchemaRegistry
+from query_search_agent import QuerySearchAgent
+import os
+
+def clean_files(registry_path, log_path):
+    for p in [registry_path, log_path]:
+        if os.path.exists(p):
+            os.remove(p)
+
+class TestSchemaTools(unittest.TestCase):
+    def setUp(self):
+        clean_files('schema_registry.json', 'schema_audit.log')
+        self.registry = SchemaRegistry('schema_registry.json', 'schema_audit.log')
+        self.agent = QuerySearchAgent(self.registry)
+        self.user_schema = {
+            'id': {'type': 'int', 'required': True},
+            'username': {'type': 'string', 'required': True},
+            'email': {'type': 'string', 'required': True},
+            'age': {'type': 'int', 'required': False}
+        }
+        self.registry.register_schema('user', self.user_schema, 'User schema for test')
+        self.short_user_schema = {
+            'username': {'type': 'string', 'required': True},
+            'age': {'type': 'int', 'required': False}
+        }
+        self.registry.register_schema('short_user', self.short_user_schema, 'Short user')
+        self.registry.register_mapping('user', 'short_user', {
+            'username': 'username',
+            'age': 'age'
+        })
+
+    def tearDown(self):
+        clean_files('schema_registry.json', 'schema_audit.log')
+
+    def test_schema_registration_and_mapping(self):
+        self.assertIn('user', self.registry.list_schemas())
+        self.assertIn('user->short_user', self.registry.list_mappings())
+
+    def test_validation(self):
+        valid_user = {'id': 1, 'username': 'alice', 'email': 'alice@email.com', 'age': 23}
+        res, errs = self.registry.validate_data('user', valid_user)
+        self.assertTrue(res)
+        self.assertEqual(errs, [])
+        missing_field = {'username': 'bob', 'email': 'bob@email.com'}
+        res, errs = self.registry.validate_data('user', missing_field)
+        self.assertFalse(res)
+        self.assertTrue(any('Required field' in err for err in errs))
+
+    def test_transform_data(self):
+        user = {'id': 1, 'username': 'alice', 'email': 'alice@email.com', 'age': 23}
+        transformed, errs = self.registry.transform_data(user, 'user', 'short_user')
+        self.assertIsNotNone(transformed)
+        self.assertEqual(set(transformed.keys()), set(['username', 'age']))
+
+    def test_search_and_query(self):
+        users = [
+            {'id': 1, 'username': 'alice', 'email': 'alice@email.com', 'age': 23},
+            {'id': 2, 'username': 'bob', 'email': 'bob@email.com', 'age': 28}
+        ]
+        results, errors = self.agent.search('user', {'username': 'alice'}, users)
+        self.assertEqual(len(results), 1)
+        self.assertEqual(results[0]['username'], 'alice')
+
+    def test_advanced_query_and_errors(self):
+        users = [
+            {'id': 1, 'username': 'alice', 'email': 'alice@email.com', 'age': 23},
+            {'id': 2, 'username': 'bob', 'email': 'bob@email.com'}
+        ]
+        adv_results, adv_errors = self.agent.advanced_query('user', 'short_user', {'username': 'bob'}, users)
+        # Should have error because of missing age but it's optional in short_user
+        self.assertEqual(len(adv_results), 1)
+        self.assertEqual(adv_results[0]['username'], 'bob')
+        # Error handling for invalid mapping
+        adv_results, adv_errors = self.agent.advanced_query('user', 'short_user', {'username': 'notfound'}, users)
+        self.assertEqual(len(adv_results), 0)
+
+    def test_json_round_trip(self):
+        obj = {'id': 42, 'name': 'Test'}
+        restored = self.agent.round_trip_json(obj)
+        self.assertEqual(obj, restored)
+
+if __name__ == '__main__':
+    unittest.main()