test_process_blob.py

import sys
import json
import os
import pandas as pd
from unittest.mock import patch, MagicMock
from uncap_app.process_blob import process_blob

def load_local_settings():
    """Loads settings from local.settings.json into the environment."""
    settings_path = os.path.join(os.path.dirname(__file__), 'local.settings.json')
    if not os.path.exists(settings_path):
        print(f"Warning: '{settings_path}' not found. Assuming environment is already configured.")
        return

    with open(settings_path, 'r') as f:
        settings = json.load(f)
    
    if 'Values' in settings:
        for key, value in settings['Values'].items():
            os.environ[key] = value

def inspect_object(obj):
    """Converts a SQLAlchemy model object to a dictionary for printing."""
    return {c.key: getattr(obj, c.key) for c in obj.__table__.columns if c.key != 'id'}

def run_test(json_data, site):
    """
    Simulates the execution of process_blob for a given JSON file, dumping the
    data that would be committed to CSV files for verification.
    """
    # These imports are moved here to ensure the environment is set before they run
    from uncap_app.core import SessionLocal
    from uncap_app.process_blob import process_blob

    mock_blob = MagicMock()
    mock_blob.name = f"{site}_test_data.json"
    mock_blob.container = "test-container"
    mock_blob.read.return_value = json.dumps(json_data).encode('utf-8')

    mock_retry_storage = {}

    def mock_get_retry_state(client, container, site, keyword):
        state_key = f"{site}/{keyword}"
        return mock_retry_storage.get(state_key, 0)

    def mock_update_retry_state(client, container, site, keyword, count):
        state_key = f"{site}/{keyword}"
        print(f"[TEST] Would update retry state for '{state_key}' to {count}.")
        mock_retry_storage[state_key] = count

    def mock_delete_retry_state(client, container, site, keyword):
        state_key = f"{site}/{keyword}"
        print(f"[TEST] Would delete retry state for '{state_key}'.")
        if state_key in mock_retry_storage:
            del mock_retry_storage[state_key]

    # The main test execution block
    with patch('uncap_app.process_blob._get_blob_service_client'), \
         patch('uncap_app.process_blob._get_retry_state', side_effect=mock_get_retry_state), \
         patch('uncap_app.process_blob._update_retry_state', side_effect=mock_update_retry_state), \
         patch('uncap_app.process_blob._delete_retry_state', side_effect=mock_delete_retry_state), \
         patch('uncap_app.process_blob.run_iceberg_batch') as mock_run_batch:

        session = SessionLocal()
        try:
            # This function will replace session.commit
            def inspect_and_dump_session():
                print("\n--- DATABASE TRANSACTION SIMULATION ---")
                data_by_table = {}

                # Capture all objects pending insertion *before* flushing
                # This includes objects added by process_blob and its sub-functions
                all_new_objects = list(session.new)

                if all_new_objects:
                    # Flush the session to ensure relationships and foreign keys are populated
                    # without committing the transaction. This is crucial for objects
                    # that need their IDs before further processing.
                    session.flush()

                    for obj in all_new_objects:
                        table_name = obj.__class__.__name__
                        if table_name not in data_by_table:
                            data_by_table[table_name] = []
                        data_by_table[table_name].append(inspect_object(obj))
                    
                    output_dir = "test_data"
                    os.makedirs(output_dir, exist_ok=True)

                    print(f"Dumping data to CSV files in '{output_dir}' directory...")
                    for table_name, records in data_by_table.items():
                        df = pd.DataFrame(records)
                        file_path = os.path.join(output_dir, f"{table_name}.csv")
                        df.to_csv(file_path, index=False)
                        print(f" -> Saved {len(records)} records to {file_path}")
                else:
                    print("No new or modified data was found in the session to commit.")

            # List to accumulate all objects added to the session
            all_added_objects = []

            # Patch session.add to intercept objects and store them
            original_session_add = session.add
            def mock_session_add(obj):
                all_added_objects.append(obj)
                original_session_add(obj) # Call the original add to maintain session state
            session.add = mock_session_add

            # This function will replace session.commit
            def inspect_and_dump_session():
                print("\n--- DATABASE TRANSACTION SIMULATION ---")
                data_by_table = {}

                # Use the accumulated list of all added objects
                if all_added_objects:
                    # Flush the session to ensure relationships and foreign keys are populated
                    # without committing the transaction. This is crucial for objects
                    # that need their IDs before further processing.
                    session.flush()

                    for obj in all_added_objects:
                        table_name = obj.__class__.__name__
                        if table_name not in data_by_table:
                            data_by_table[table_name] = []
                        data_by_table[table_name].append(inspect_object(obj))
                    
                    output_dir = "test_data"
                    os.makedirs(output_dir, exist_ok=True)

                    print(f"Dumping data to CSV files in '{output_dir}' directory...")
                    for table_name, records in data_by_table.items():
                        df = pd.DataFrame(records)
                        file_path = os.path.join(output_dir, f"{table_name}.csv")
                        df.to_csv(file_path, index=False)
                        print(f" -> Saved {len(records)} records to {file_path}")
                else:
                    print("No new or modified data was found in the session to commit.")

            session.commit = inspect_and_dump_session

            # Ensure process_blob uses our modified session
            with patch('uncap_app.process_blob.SessionLocal', return_value=session):
                process_blob(mock_blob, "test-container", site, min_skus_per_keyword=1)

            print("\n--- ICEBERG API CALL SIMULATION ---")
            if mock_run_batch.called:
                call_args = mock_run_batch.call_args[0]
                print(f"run_iceberg_batch would have been called for site '{call_args[0]}' with keywords:")
                print(json.dumps(call_args[1], indent=2))
            else:
                print("run_iceberg_batch was not called.")

        finally:
            session.rollback() # Ensure the real DB is never touched
            session.close()

if __name__ == "__main__":
    load_local_settings()
    sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

    if len(sys.argv) != 2:
        print(f"Usage: python {sys.argv[0]} <path_to_json_file>", file=sys.stderr)
        sys.exit(1)

    file_path = sys.argv[1]
    site = "wayfair"

    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        run_test(data, site)

    except FileNotFoundError:
        print(f"Error: File not found at {file_path}", file=sys.stderr)
        sys.exit(1)
    except json.JSONDecodeError:
        print(f"Error: Could not decode JSON from {file_path}", file=sys.stderr)
        sys.exit(1)