Spaces:
Sleeping
Sleeping
File size: 4,263 Bytes
765341b 15bb604 765341b 15bb604 765341b 15bb604 765341b 15bb604 765341b 15bb604 765341b 15bb604 765341b 15bb604 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """
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()
|