Spaces:
Sleeping
Sleeping
| """ | |
| Test script for batch missing data recovery, validation, and lead status flows | |
| """ | |
| 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) | |
| # Add new columns for data recovery & validation | |
| self.user_schema = { | |
| 'id': {'type': 'int', 'required': True}, | |
| 'username': {'type': 'string', 'required': True}, | |
| 'email': {'type': 'string', 'required': True}, | |
| 'age': {'type': 'int', 'required': False}, | |
| 'missing_data_flag': {'type': 'bool', 'required': False, 'default': False}, | |
| 'data_recovery_tag': {'type': 'string', 'required': False, 'default': ''}, | |
| 'recovery_attempted': {'type': 'bool', 'required': False, 'default': False}, | |
| 'lead_status': {'type': 'string', 'required': False, 'default': 'active'}, | |
| 'scrub_reason': {'type': 'string', 'required': False, 'default': ''}, | |
| } | |
| self.registry.register_schema('user', self.user_schema, 'User schema with recovery fields') | |
| def test_batch_missing_data_and_recovery(self): | |
| # Simulating a batch of leads, some missing required fields | |
| leads = [ | |
| {'id': 1, 'username': 'alice', 'email': 'alice@email.com', 'age': 23}, | |
| {'id': 2, 'username': 'bob', 'email': '', 'age': 31}, # Missing email | |
| {'id': 3, 'username': '', 'email': 'sue@email.com', 'age': 22}, # Missing username | |
| {'id': 4, 'username': 'ed', 'email': 'ed@email.com'}, # Partial | |
| ] | |
| processed = [] | |
| recovery_queue = [] | |
| for lead in leads: | |
| res, errs = self.registry.validate_data('user', lead) | |
| if res: | |
| lead.update({'missing_data_flag': False, 'data_recovery_tag': '', 'recovery_attempted': False, 'lead_status': 'active', 'scrub_reason': ''}) | |
| else: | |
| lead['missing_data_flag'] = True | |
| lead['data_recovery_tag'] = 'RECOVERY_TRIGGERED' | |
| lead['recovery_attempted'] = False | |
| lead['lead_status'] = 'pending_recovery' | |
| lead['scrub_reason'] = '' | |
| recovery_queue.append(lead) | |
| processed.append(lead) | |
| self.assertEqual(len(recovery_queue), 3) | |
| # Batch: Mark recovery attempted & try recovery (simulate only 1 can be rescued) | |
| for lead in recovery_queue: | |
| lead['recovery_attempted'] = True | |
| # Try recovery - for the test, rescue only id==2 (fill email) | |
| if lead['id'] == 2: | |
| lead['email'] = 'bob@email.com' | |
| lead['missing_data_flag'] = False | |
| lead['data_recovery_tag'] = '' | |
| lead['lead_status'] = 'active' | |
| else: | |
| lead['missing_data_flag'] = True | |
| lead['lead_status'] = 'needs_human_review' | |
| recovery_for_human = [l for l in recovery_queue if l['lead_status'] == 'needs_human_review'] | |
| self.assertEqual(len(recovery_for_human), 2) | |
| # Human reviews, decides to approve id==3 and scrub id==4 | |
| for l in recovery_for_human: | |
| if l['id'] == 3: | |
| l['lead_status'] = 'active' # Approved back | |
| l['scrub_reason'] = '' | |
| if l['id'] == 4: | |
| l['lead_status'] = 'scrubbed' | |
| l['scrub_reason'] = 'No username, unresolvable after recovery attempt.' | |
| scrubbed_leads = [l for l in processed if l['lead_status'] == 'scrubbed'] | |
| self.assertTrue(any(l['id']==4 for l in scrubbed_leads)) | |
| # All others are either reactivated or pending normal | |
| for l in processed: | |
| self.assertIn(l['lead_status'], ['active','scrubbed','needs_human_review','pending_recovery']) | |
| if __name__ == '__main__': | |
| unittest.main() | |