Task 3. Phân phối sự kiện và cấu trúc Kafka Topics

# Setup path and load configurations
import os
import sys
import yaml
import subprocess
from pathlib import Path

def find_project_root() -> Path:
    p = Path(os.getcwd()).resolve()
    for parent in [p] + list(p.parents):
        if (parent / '.env').exists() or (parent / 'pyproject.toml').exists():
            return parent
    return p

project_root = find_project_root()
os.chdir(str(project_root))
sys.path.append(str(project_root / 'src'))
sys.path.append(str(project_root / 'scripts'))

# Verify Kafka Docker container status
res = subprocess.run(['docker', 'compose', '--env-file', '.env', '-f', 'infra/docker-compose.yml', 'ps', 'kafka', '--format', 'json'], capture_output=True, text=True)
print('Kafka container state:', 'RUNNING' if 'running' in res.stdout.lower() else 'STOPPED')
assert 'running' in res.stdout.lower(), 'Kafka container is not running'

# Read target partitions config from config/topics.yaml
config_path = project_root / 'config' / 'topics.yaml'
with open(config_path, 'r') as f:
    topics_config = yaml.safe_load(f) or {}
print('Configured topics:', [t['name'] for t in topics_config.get('topics', [])])

# Fetch topic partition counts using Kafka GetOffsetShell via helper
from infrastructure.verification.kafka_connect import get_topic_end_offsets
bootstrap_servers = 'localhost:9092'
expected_topics = {'cpg.nodes': 3, 'cpg.edges': 3, 'source.metadata': 1, 'parser.errors': 1}

print(f'{"Topic":<20} | {"Partitions":<10} | {"Status":<8}')
print('-' * 45)
for topic, expected_partitions in expected_topics.items():
    try:
        offsets = get_topic_end_offsets(bootstrap_servers, topic)
        p_count = len(offsets)
        print(f'{topic:<20} | {p_count:<10} | {"OK [PASS]":<8}')
        assert p_count == expected_partitions, f'Topic {topic} partition count mismatch: {p_count} vs {expected_partitions}'
    except Exception as exc:
        print(f'{topic:<20} | {"ERROR":<10} | {str(exc):<8}')
        raise exc

# Execute fresh parse on the target python file using an isolated SQLite state database
import json
from confluent_kafka import Consumer, TopicPartition
from infrastructure.verification.kafka_connect import get_topic_end_offsets

state_db = 'workspace/tmp/task3-notebook/state.sqlite3'
os.makedirs('workspace/tmp/task3-notebook', exist_ok=True)
if os.path.exists(state_db):
    os.remove(state_db)

target_file = '.github/scripts/assign_reviewers.py'
# Capture before offsets
topics = ['cpg.nodes', 'cpg.edges', 'source.metadata', 'parser.errors']
before_offsets = {t: get_topic_end_offsets(bootstrap_servers, t) for t in topics}

# Run CLI Parser Service
cmd = [
    'uv', 'run', 'lab04', 'parse-file',
    '--file', target_file,
    '--no-dry-run'
]
env_override = dict(os.environ, PARSER_STATE_DB=state_db)
res_parse = subprocess.run(cmd, env=env_override, capture_output=True, text=True)
print('Parser CLI status:', 'SUCCESS' if res_parse.returncode == 0 else 'FAILED')
if res_parse.returncode != 0:
    print('Error output:', res_parse.stderr)
assert res_parse.returncode == 0, 'Parser Service execution failed'

# Capture after offsets
after_offsets = {t: get_topic_end_offsets(bootstrap_servers, t) for t in topics}
deltas = {}
for t in topics:
    deltas[t] = {p: after_offsets[t].get(p, 0) - before_offsets[t].get(p, 0) for p in after_offsets[t]}
    print(f'Topic {t} published events per partition: {deltas[t]}')

# Consume published messages and validate key, schema, and routing
from infrastructure.messaging.event_validator import EventValidator
validator = EventValidator(schemas_dir=Path('schemas'))

conf = {
    'bootstrap.servers': bootstrap_servers,
    'group.id': 'task3-notebook-validator-group',
    'auto.offset.reset': 'earliest',
    'enable.auto.commit': False
}
consumer = Consumer(conf)
samples = {}

for topic in topics:
    t_deltas = deltas[topic]
    for partition, count in t_deltas.items():
        if count == 0:
            continue
        tp = TopicPartition(topic, partition, before_offsets[topic][partition])
        consumer.assign([tp])
        for _ in range(count):
            msg = consumer.poll(timeout=5.0)
            assert msg is not None, f'Failed to poll message from {topic}-{partition}'
            assert msg.error() is None, f'Consumer error: {msg.error()}'

            # Check key is string and equals file_id
            key_bytes = msg.key()
            assert key_bytes is not None, 'Kafka message key is missing'
            key_str = key_bytes.decode('utf-8')

            # Parse value
            val_bytes = msg.value()
            assert val_bytes is not None, 'Kafka message value is missing'
            val_json = json.loads(val_bytes.decode('utf-8'))

            # Validate using schema
            validator.validate(val_json.get('event_type'), val_json)
            assert val_json.get('file_id') == key_str, f'Key {key_str} mismatch with event file_id'

            # Store sample
            if topic not in samples:
                samples[topic] = val_json
consumer.close()
print('Verification successful: Keys and schemas are valid!')

# Display a clean sample event for each populated topic
for topic, payload in samples.items():
    print(f'=== Sample Event for Topic: {topic} ===')
    truncated_payload = {
        'schema_version': payload.get('schema_version'),
        'event_id': payload.get('event_id'),
        'event_type': payload.get('event_type'),
        'event_time': payload.get('event_time'),
        'file_id': payload.get('file_id'),
        'file_path': payload.get('file_path'),
        'content_hash': payload.get('content_hash'),
        'parser_version': payload.get('parser_version'),
    }
    print(json.dumps(truncated_payload, indent=2))
    print('-' * 40)

# Run unchanged execution and capture topic deltas
before_offsets_rerun = {t: get_topic_end_offsets(bootstrap_servers, t) for t in topics}

res_parse_rerun = subprocess.run(cmd, env=env_override, capture_output=True, text=True)
print('Rerun stdout status:', 'SUCCESS' if res_parse_rerun.returncode == 0 else 'FAILED')
assert res_parse_rerun.returncode == 0, 'Rerun failed'

# Verify from stdout that it skipped the file
print('Output:', res_parse_rerun.stdout.strip())

after_offsets_rerun = {t: get_topic_end_offsets(bootstrap_servers, t) for t in topics}
for t in topics:
    delta = sum(after_offsets_rerun[t].get(p, 0) - before_offsets_rerun[t].get(p, 0) for p in after_offsets_rerun[t])
    print(f'Topic {t} delta: {delta}')
    assert delta == 0, f'Expected 0 messages on {t}, got {delta}'
print('[PASS] Unchanged file skip verification completed successfully.')

# Write syntax error file, run parser, assert event routing to parser.errors
temp_err_file = Path('workspace/source/transformers-pr-agent/error_file.py')
temp_err_file.write_text('x = \n', encoding='utf-8')

before_offsets_err = {t: get_topic_end_offsets(bootstrap_servers, t) for t in topics}

cmd_err = [
    'uv', 'run', 'lab04', 'parse-file',
    '--file', 'error_file.py',
    '--no-dry-run'
]
res_parse_err = subprocess.run(cmd_err, env=env_override, capture_output=True, text=True)
# Expect exit code 1 due to syntax error
print('Parser CLI exit code:', res_parse_err.returncode)
assert res_parse_err.returncode != 0, 'Expected failure exit code, but succeeded'

after_offsets_err = {t: get_topic_end_offsets(bootstrap_servers, t) for t in topics}
deltas_err = {t: sum(after_offsets_err[t].get(p, 0) - before_offsets_err[t].get(p, 0) for p in after_offsets_err[t]) for t in topics}
print('Deltas:', deltas_err)

assert deltas_err['parser.errors'] == 1, f'Expected 1 event on parser.errors, got {deltas_err["parser.errors"]}'
assert deltas_err['cpg.nodes'] == 0, 'Expected 0 node events'
assert deltas_err['cpg.edges'] == 0, 'Expected 0 edge events'
assert deltas_err['source.metadata'] == 0, 'Expected 0 metadata events'

# Consume the error event to validate key and schema
conf_err = dict(conf)
conf_err['group.id'] = 'task3-notebook-err-group'
consumer_err = Consumer(conf_err)
tp_err = TopicPartition('parser.errors', 0, before_offsets_err['parser.errors'][0])
consumer_err.assign([tp_err])
msg_err = consumer_err.poll(timeout=5.0)
assert msg_err is not None, 'Failed to poll error event'
assert msg_err.error() is None

key_bytes = msg_err.key()
assert key_bytes is not None, 'Kafka message key is missing'
key_str = key_bytes.decode('utf-8')

val_json = json.loads(msg_err.value().decode('utf-8'))
validator.validate(val_json.get('event_type'), val_json)
assert val_json.get('file_id') == key_str, 'Key mismatch'
consumer_err.close()

# Clean up temporary file
if temp_err_file.exists():
    temp_err_file.unlink()
print('[PASS] Parser error verification completed successfully.')