diff --git a/run_tests.py b/run_tests.py
index b59a728..04b920c 100644
--- a/run_tests.py
+++ b/run_tests.py
@@ -961,31 +961,31 @@ def run_fallback_report() -> TestResult:
     "file_upload",
     ["tests/test_file_upload_comprehensive.py"],
     "Running comprehensive file upload testing suite",
-    "file_upload"
+    "file_upload",
 )
 
 NEW_SUITES.add_test(
-    "error_boundary", 
+    "error_boundary",
     ["tests/test_error_boundary_comprehensive.py"],
     "Running comprehensive error boundary testing suite",
-    "error_boundary"
+    "error_boundary",
 )
 
 NEW_SUITES.add_test(
     "visualization",
-    ["tests/test_visualization_pipeline_comprehensive.py"], 
+    ["tests/test_visualization_pipeline_comprehensive.py"],
     "Running comprehensive visualization pipeline testing suite",
-    "visualization"
+    "visualization",
 )
 
 NEW_SUITES.add_test(
     "all_new_suites",
     [
         "tests/test_file_upload_comprehensive.py",
-        "tests/test_error_boundary_comprehensive.py", 
-        "tests/test_visualization_pipeline_comprehensive.py"
+        "tests/test_error_boundary_comprehensive.py",
+        "tests/test_visualization_pipeline_comprehensive.py",
     ],
-    "Running all new comprehensive testing suites"
+    "Running all new comprehensive testing suites",
 )
 
 # Update the TEST_CATEGORIES to include process mining, UI tests, and new suites
@@ -1286,7 +1286,9 @@ def main():
         "--process-mining-all", action="store_true", help="Run all process mining tests"
     )
     parser.add_argument("--ui-all", action="store_true", help="Run all UI tests")
-    parser.add_argument("--new-suites-all", action="store_true", help="Run all new comprehensive test suites")
+    parser.add_argument(
+        "--new-suites-all", action="store_true", help="Run all new comprehensive test suites"
+    )
 
     # Category specific test options
     parser.add_argument(
diff --git a/tests/test_app_ui.py b/tests/test_app_ui.py
index 96c97df..95e6d20 100644
--- a/tests/test_app_ui.py
+++ b/tests/test_app_ui.py
@@ -37,9 +37,26 @@ def load_sample_data(self) -> None:
         This helper abstracts the data loading flow, making tests immune to
         changes in the data loading implementation.
         """
+        import time
+
         try:
             # Click the Load Sample Data button using its stable key with increased timeout
-            self.at.button(key="load_sample_data_button").click().run(timeout=10)
+            self.at.button(key="load_sample_data_button").click().run(timeout=15)
+
+            # Wait for data loading to complete by checking session state
+            max_retries = 5
+            for attempt in range(max_retries):
+                if (
+                    hasattr(self.at.session_state, "events_data")
+                    and self.at.session_state.events_data is not None
+                    and len(self.at.session_state.events_data) > 0
+                ):
+                    # Data loaded successfully, now wait for event_statistics
+                    break
+                time.sleep(0.5)  # Wait 500ms between retries
+                if attempt < max_retries - 1:  # Don't run on last attempt
+                    self.at.run(timeout=10)  # Re-run to refresh state
+
         except Exception as e:
             # If button interaction fails, manually load sample data for testing
             from datetime import datetime, timedelta
@@ -67,10 +84,52 @@ def load_sample_data(self) -> None:
 
             self.at.session_state.events_data = pd.DataFrame(data)
 
+            # CRITICAL FIX: Calculate event_statistics manually when using fallback
+            # Import the function from app.py
+            import os
+            import sys
+
+            sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+            from app import get_event_statistics
+
+            self.at.session_state.event_statistics = get_event_statistics(
+                self.at.session_state.events_data
+            )
+
         # Verify data was loaded by checking session state
         assert self.at.session_state.events_data is not None, "Sample data should be loaded"
         assert len(self.at.session_state.events_data) > 0, "Sample data should not be empty"
 
+        # Wait for event_statistics to be calculated (give it more time)
+        max_retries = 10
+        for attempt in range(max_retries):
+            if (
+                hasattr(self.at.session_state, "event_statistics")
+                and len(self.at.session_state.event_statistics) > 0
+            ):
+                break
+            time.sleep(0.3)  # Wait 300ms between retries
+            if attempt < max_retries - 1:  # Don't run on last attempt
+                self.at.run(
+                    timeout=10
+                )  # Re-run to refresh state and trigger statistics calculation
+
+        # Final verification that event_statistics was calculated
+        if (
+            not hasattr(self.at.session_state, "event_statistics")
+            or len(self.at.session_state.event_statistics) == 0
+        ):
+            # Last resort: manually calculate if still not present
+            import os
+            import sys
+
+            sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+            from app import get_event_statistics
+
+            self.at.session_state.event_statistics = get_event_statistics(
+                self.at.session_state.events_data
+            )
+
     def build_funnel(self, steps: List[str]) -> None:
         """
         Build a funnel by selecting the specified steps.
diff --git a/tests/test_app_ui_phase1_golden.py b/tests/test_app_ui_phase1_golden.py
index 9c52fd1..5420b0f 100644
--- a/tests/test_app_ui_phase1_golden.py
+++ b/tests/test_app_ui_phase1_golden.py
@@ -12,9 +12,38 @@
 4. Cover the complete "happy path" and key edge cases
 """
 
+import time
+
 from streamlit.testing.v1 import AppTest
 
 APP_FILE = "app.py"
+DEFAULT_TIMEOUT = 20  # Increased from 10 to handle slower module loading in Python 3.12
+
+
+def wait_for_condition(at, condition_func, max_retries=10, sleep_time=0.5, timeout=15):
+    """
+    Robust waiting mechanism for test conditions.
+
+    Args:
+        at: AppTest instance
+        condition_func: Function that returns True when condition is met
+        max_retries: Maximum number of retries
+        sleep_time: Time to sleep between retries
+        timeout: Timeout for each at.run() call
+    """
+    for attempt in range(max_retries):
+        try:
+            if condition_func():
+                return True
+        except (AttributeError, KeyError, TypeError):
+            # Condition not ready yet, continue waiting
+            pass
+
+        if attempt < max_retries - 1:
+            time.sleep(sleep_time)
+            at.run(timeout=timeout)
+
+    return False
 
 
 class TestPhase1GoldenStandard:
@@ -30,10 +59,23 @@ def test_full_successful_analysis_flow(self):
         This test validates the full workflow from data loading to analysis results.
         It must pass before, during, and after Phase 1 refactoring.
         """
-        at = AppTest.from_file(APP_FILE, default_timeout=10).run()
-
-        # Step 1: Load sample data
-        at.button(key="load_sample_data_button").click().run()
+        at = AppTest.from_file(APP_FILE, default_timeout=DEFAULT_TIMEOUT).run()
+
+        # Step 1: Load sample data with robust waiting
+        at.button(key="load_sample_data_button").click().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for data to be fully loaded and processed
+        data_loaded = wait_for_condition(
+            at,
+            lambda: (
+                hasattr(at.session_state, "events_data")
+                and at.session_state.events_data is not None
+                and len(at.session_state.events_data) > 0
+            ),
+            max_retries=15,
+            timeout=DEFAULT_TIMEOUT,
+        )
+        assert data_loaded, "Events data should be loaded within timeout"
 
         # Verify data was loaded (check session state)
         assert at.session_state.events_data is not None, "Events data should be loaded"
@@ -48,7 +90,13 @@ def test_full_successful_analysis_flow(self):
 
         for event in selected_events:
             checkbox_key = f"event_cb_{event.replace(' ', '_').replace('-', '_')}"
-            at.checkbox(key=checkbox_key).check().run()
+            at.checkbox(key=checkbox_key).check().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for funnel steps to be updated
+        funnel_built = wait_for_condition(
+            at, lambda: len(at.session_state.funnel_steps) == 3, timeout=DEFAULT_TIMEOUT
+        )
+        assert funnel_built, "Funnel should be built within timeout"
 
         # Verify funnel steps were added to session state
         assert len(at.session_state.funnel_steps) == 3, "Should have 3 funnel steps"
@@ -56,8 +104,20 @@ def test_full_successful_analysis_flow(self):
             "Steps should match selected events"
         )
 
-        # Step 3: Run analysis
-        at.button(key="analyze_funnel_button").click().run()
+        # Step 3: Run analysis with robust waiting
+        at.button(key="analyze_funnel_button").click().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for analysis to complete
+        analysis_complete = wait_for_condition(
+            at,
+            lambda: (
+                hasattr(at.session_state, "analysis_results")
+                and at.session_state.analysis_results is not None
+            ),
+            max_retries=20,  # Analysis can take longer
+            timeout=DEFAULT_TIMEOUT,
+        )
+        assert analysis_complete, "Analysis should complete within timeout"
 
         # Verify analysis results were generated
         assert at.session_state.analysis_results is not None, (
@@ -76,23 +136,47 @@ def test_state_reset_on_clear_all(self):
         This test ensures that the Clear All button properly resets the funnel state
         and that the UI reflects this change correctly.
         """
-        at = AppTest.from_file(APP_FILE, default_timeout=10).run()
-
-        # Load data and build funnel
-        at.button(key="load_sample_data_button").click().run()
+        at = AppTest.from_file(APP_FILE, default_timeout=DEFAULT_TIMEOUT).run()
+
+        # Load data and build funnel with robust waiting
+        at.button(key="load_sample_data_button").click().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for data loading
+        data_loaded = wait_for_condition(
+            at,
+            lambda: (
+                hasattr(at.session_state, "events_data")
+                and at.session_state.events_data is not None
+                and len(at.session_state.events_data) > 0
+            ),
+            timeout=DEFAULT_TIMEOUT,
+        )
+        assert data_loaded, "Data should be loaded"
 
         available_events = sorted(at.session_state.events_data["event_name"].unique())
         first_event = available_events[0]
         checkbox_key = f"event_cb_{first_event.replace(' ', '_').replace('-', '_')}"
 
-        at.checkbox(key=checkbox_key).check().run()
+        at.checkbox(key=checkbox_key).check().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for step to be added
+        step_added = wait_for_condition(
+            at, lambda: len(at.session_state.funnel_steps) == 1, timeout=DEFAULT_TIMEOUT
+        )
+        assert step_added, "Step should be added"
 
         # Verify step was added
         assert len(at.session_state.funnel_steps) == 1, "Should have 1 funnel step"
         assert first_event in at.session_state.funnel_steps, "Event should be in funnel steps"
 
         # Clear all steps
-        at.button(key="clear_all_button").click().run()
+        at.button(key="clear_all_button").click().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for state to be cleared
+        state_cleared = wait_for_condition(
+            at, lambda: len(at.session_state.funnel_steps) == 0, timeout=DEFAULT_TIMEOUT
+        )
+        assert state_cleared, "State should be cleared"
 
         # Verify state was cleared
         assert len(at.session_state.funnel_steps) == 0, "Funnel steps should be empty after clear"
@@ -108,24 +192,48 @@ def test_event_selection_state_management(self):
         This test verifies that event selection/deselection properly updates
         the session state and that the UI reflects these changes.
         """
-        at = AppTest.from_file(APP_FILE, default_timeout=10).run()
-
-        # Load data
-        at.button(key="load_sample_data_button").click().run()
+        at = AppTest.from_file(APP_FILE, default_timeout=DEFAULT_TIMEOUT).run()
+
+        # Load data with robust waiting
+        at.button(key="load_sample_data_button").click().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for data loading
+        data_loaded = wait_for_condition(
+            at,
+            lambda: (
+                hasattr(at.session_state, "events_data")
+                and at.session_state.events_data is not None
+                and len(at.session_state.events_data) > 0
+            ),
+            timeout=DEFAULT_TIMEOUT,
+        )
+        assert data_loaded, "Data should be loaded"
 
         available_events = sorted(at.session_state.events_data["event_name"].unique())
         test_events = available_events[:2]
 
         # Select first event
         checkbox_key_1 = f"event_cb_{test_events[0].replace(' ', '_').replace('-', '_')}"
-        at.checkbox(key=checkbox_key_1).check().run()
+        at.checkbox(key=checkbox_key_1).check().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for first event selection
+        first_selected = wait_for_condition(
+            at, lambda: len(at.session_state.funnel_steps) == 1, timeout=DEFAULT_TIMEOUT
+        )
+        assert first_selected, "First event should be selected"
 
         assert len(at.session_state.funnel_steps) == 1, "Should have 1 step after first selection"
         assert test_events[0] in at.session_state.funnel_steps, "First event should be selected"
 
         # Select second event
         checkbox_key_2 = f"event_cb_{test_events[1].replace(' ', '_').replace('-', '_')}"
-        at.checkbox(key=checkbox_key_2).check().run()
+        at.checkbox(key=checkbox_key_2).check().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for second event selection
+        second_selected = wait_for_condition(
+            at, lambda: len(at.session_state.funnel_steps) == 2, timeout=DEFAULT_TIMEOUT
+        )
+        assert second_selected, "Second event should be selected"
 
         assert len(at.session_state.funnel_steps) == 2, (
             "Should have 2 steps after second selection"
@@ -133,7 +241,18 @@ def test_event_selection_state_management(self):
         assert test_events[1] in at.session_state.funnel_steps, "Second event should be selected"
 
         # Deselect first event
-        at.checkbox(key=checkbox_key_1).uncheck().run()
+        at.checkbox(key=checkbox_key_1).uncheck().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for first event deselection
+        first_deselected = wait_for_condition(
+            at,
+            lambda: (
+                len(at.session_state.funnel_steps) == 1
+                and test_events[0] not in at.session_state.funnel_steps
+            ),
+            timeout=DEFAULT_TIMEOUT,
+        )
+        assert first_deselected, "First event should be deselected"
 
         assert len(at.session_state.funnel_steps) == 1, "Should have 1 step after deselection"
         assert test_events[0] not in at.session_state.funnel_steps, (
@@ -153,14 +272,27 @@ def test_data_loading_initializes_session_state(self):
         This test ensures that loading data sets up all the necessary session state
         variables that other components depend on.
         """
-        at = AppTest.from_file(APP_FILE, default_timeout=10).run()
+        at = AppTest.from_file(APP_FILE, default_timeout=DEFAULT_TIMEOUT).run()
 
         # Verify initial state
         assert at.session_state.events_data is None, "Events data should be None initially"
         assert len(at.session_state.funnel_steps) == 0, "Funnel steps should be empty initially"
 
-        # Load data
-        at.button(key="load_sample_data_button").click().run()
+        # Load data with robust waiting
+        at.button(key="load_sample_data_button").click().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for data loading to complete
+        data_loaded = wait_for_condition(
+            at,
+            lambda: (
+                hasattr(at.session_state, "events_data")
+                and at.session_state.events_data is not None
+                and len(at.session_state.events_data) > 0
+            ),
+            max_retries=15,
+            timeout=DEFAULT_TIMEOUT,
+        )
+        assert data_loaded, "Events data should be loaded within timeout"
 
         # Verify data loading initialized required state
         assert at.session_state.events_data is not None, "Events data should be loaded"
@@ -171,6 +303,18 @@ def test_data_loading_initializes_session_state(self):
         for col in required_columns:
             assert col in at.session_state.events_data.columns, f"Column {col} should exist"
 
+        # Wait for event statistics to be generated
+        stats_generated = wait_for_condition(
+            at,
+            lambda: (
+                hasattr(at.session_state, "event_statistics")
+                and len(at.session_state.event_statistics) > 0
+            ),
+            max_retries=10,
+            timeout=DEFAULT_TIMEOUT,
+        )
+        assert stats_generated, "Event statistics should be generated within timeout"
+
         # Verify event statistics were generated
         assert hasattr(at.session_state, "event_statistics"), (
             "Event statistics should be generated"
@@ -189,10 +333,22 @@ def test_analysis_requires_minimum_steps(self):
         This test ensures that the analysis function properly validates
         that at least 2 steps are required for funnel analysis.
         """
-        at = AppTest.from_file(APP_FILE, default_timeout=10).run()
-
-        # Load data
-        at.button(key="load_sample_data_button").click().run()
+        at = AppTest.from_file(APP_FILE, default_timeout=DEFAULT_TIMEOUT).run()
+
+        # Load data with robust waiting
+        at.button(key="load_sample_data_button").click().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for data loading
+        data_loaded = wait_for_condition(
+            at,
+            lambda: (
+                hasattr(at.session_state, "events_data")
+                and at.session_state.events_data is not None
+                and len(at.session_state.events_data) > 0
+            ),
+            timeout=DEFAULT_TIMEOUT,
+        )
+        assert data_loaded, "Data should be loaded"
 
         # Note: analyze_funnel_button only appears when there are funnel steps
         # So we can't test clicking it with 0 steps - the button won't exist
@@ -201,10 +357,20 @@ def test_analysis_requires_minimum_steps(self):
         available_events = sorted(at.session_state.events_data["event_name"].unique())
         first_event = available_events[0]
         checkbox_key = f"event_cb_{first_event.replace(' ', '_').replace('-', '_')}"
-        at.checkbox(key=checkbox_key).check().run()
+        at.checkbox(key=checkbox_key).check().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for step to be added
+        step_added = wait_for_condition(
+            at, lambda: len(at.session_state.funnel_steps) == 1, timeout=DEFAULT_TIMEOUT
+        )
+        assert step_added, "Step should be added"
 
         # Now the analyze button should appear, try to analyze with 1 step
-        at.button(key="analyze_funnel_button").click().run()
+        at.button(key="analyze_funnel_button").click().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait a moment for any potential analysis attempt
+        time.sleep(1)
+        at.run(timeout=DEFAULT_TIMEOUT)
 
         # Should not generate analysis results with insufficient steps (< 2)
         assert at.session_state.analysis_results is None, "Should not generate results with 1 step"
@@ -212,10 +378,25 @@ def test_analysis_requires_minimum_steps(self):
         # Add second step
         second_event = available_events[1]
         checkbox_key_2 = f"event_cb_{second_event.replace(' ', '_').replace('-', '_')}"
-        at.checkbox(key=checkbox_key_2).check().run()
+        at.checkbox(key=checkbox_key_2).check().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for second step
+        second_step_added = wait_for_condition(
+            at, lambda: len(at.session_state.funnel_steps) == 2, timeout=DEFAULT_TIMEOUT
+        )
+        assert second_step_added, "Second step should be added"
 
         # Now analysis should work with 2+ steps
-        at.button(key="analyze_funnel_button").click().run()
+        at.button(key="analyze_funnel_button").click().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for analysis to complete
+        analysis_complete = wait_for_condition(
+            at,
+            lambda: at.session_state.analysis_results is not None,
+            max_retries=20,
+            timeout=DEFAULT_TIMEOUT,
+        )
+        assert analysis_complete, "Analysis should complete with 2+ steps"
 
         # Should generate analysis results with 2+ steps
         assert at.session_state.analysis_results is not None, (
@@ -232,33 +413,75 @@ def test_session_state_persistence_across_interactions(self):
         This test ensures that session state maintains consistency across
         multiple user interactions and UI updates.
         """
-        at = AppTest.from_file(APP_FILE, default_timeout=10).run()
+        at = AppTest.from_file(APP_FILE, default_timeout=DEFAULT_TIMEOUT).run()
+
+        # Load data with robust waiting
+        at.button(key="load_sample_data_button").click().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for data loading
+        data_loaded = wait_for_condition(
+            at,
+            lambda: (
+                hasattr(at.session_state, "events_data")
+                and at.session_state.events_data is not None
+                and len(at.session_state.events_data) > 0
+            ),
+            timeout=DEFAULT_TIMEOUT,
+        )
+        assert data_loaded, "Data should be loaded"
 
-        # Load data
-        at.button(key="load_sample_data_button").click().run()
         original_data_length = len(at.session_state.events_data)
 
-        # Build funnel
+        # Build funnel with robust waiting
         available_events = sorted(at.session_state.events_data["event_name"].unique())
         selected_events = available_events[:3]
 
-        for event in selected_events:
+        for i, event in enumerate(selected_events):
             checkbox_key = f"event_cb_{event.replace(' ', '_').replace('-', '_')}"
-            at.checkbox(key=checkbox_key).check().run()
+            at.checkbox(key=checkbox_key).check().run(timeout=DEFAULT_TIMEOUT)
+
+            # Wait for each step to be added
+            step_count = i + 1
+            step_added = wait_for_condition(
+                at,
+                lambda: len(at.session_state.funnel_steps) == step_count,
+                timeout=DEFAULT_TIMEOUT,
+            )
+            assert step_added, f"Step {step_count} should be added"
 
-        # Run analysis
-        at.button(key="analyze_funnel_button").click().run()
+        # Run analysis with robust waiting
+        at.button(key="analyze_funnel_button").click().run(timeout=DEFAULT_TIMEOUT)
+
+        # Wait for analysis to complete
+        analysis_complete = wait_for_condition(
+            at,
+            lambda: at.session_state.analysis_results is not None,
+            max_retries=20,
+            timeout=DEFAULT_TIMEOUT,
+        )
+        assert analysis_complete, "Analysis should complete"
 
         # Verify all state is still consistent
         assert len(at.session_state.events_data) == original_data_length, "Data should persist"
         assert len(at.session_state.funnel_steps) == 3, "Funnel steps should persist"
         assert at.session_state.analysis_results is not None, "Analysis results should persist"
 
-        # Add another step
+        # Add another step if available
         if len(available_events) > 3:
             fourth_event = available_events[3]
             checkbox_key_4 = f"event_cb_{fourth_event.replace(' ', '_').replace('-', '_')}"
-            at.checkbox(key=checkbox_key_4).check().run()
+            at.checkbox(key=checkbox_key_4).check().run(timeout=DEFAULT_TIMEOUT)
+
+            # Wait for fourth step and results clearing
+            fourth_step_added = wait_for_condition(
+                at,
+                lambda: (
+                    len(at.session_state.funnel_steps) == 4
+                    and at.session_state.analysis_results is None
+                ),
+                timeout=DEFAULT_TIMEOUT,
+            )
+            assert fourth_step_added, "Fourth step should be added and results cleared"
 
             # Verify analysis results were cleared (as expected when funnel changes)
             assert at.session_state.analysis_results is None, (
diff --git a/tests/test_error_boundary_comprehensive.py b/tests/test_error_boundary_comprehensive.py
index dcfdd6d..dbec1cc 100644
--- a/tests/test_error_boundary_comprehensive.py
+++ b/tests/test_error_boundary_comprehensive.py
@@ -8,17 +8,14 @@
 - Graceful degradation patterns
 """
 
-import json
-import logging
 from datetime import datetime, timedelta
-from io import StringIO
 from unittest.mock import Mock, patch
 
 import pandas as pd
 import pytest
 
-from core import DataSourceManager, FunnelCalculator, FunnelConfigManager
-from models import CountingMethod, FunnelConfig, FunnelOrder, ReentryMode
+from core import DataSourceManager, FunnelCalculator
+from models import FunnelConfig
 
 
 @pytest.mark.error_boundary
@@ -52,8 +49,12 @@ def test_file_loading_exception_handling(self, data_manager):
                 result_df = data_manager.load_from_file(mock_file)
 
                 # Should return empty DataFrame
-                assert isinstance(result_df, pd.DataFrame), f"Should return DataFrame for {type(exception).__name__}"
-                assert len(result_df) == 0, f"Should return empty DataFrame for {type(exception).__name__}"
+                assert isinstance(result_df, pd.DataFrame), (
+                    f"Should return DataFrame for {type(exception).__name__}"
+                )
+                assert len(result_df) == 0, (
+                    f"Should return empty DataFrame for {type(exception).__name__}"
+                )
 
         print("✅ File loading exception handling test passed")
 
@@ -66,7 +67,6 @@ def test_data_validation_error_messages(self, data_manager):
                 pd.DataFrame({"wrong_column": [1, 2, 3]}),
                 ["missing", "required", "columns"],
             ),
-            
             # Empty DataFrame scenario
             (
                 pd.DataFrame(),
@@ -83,7 +83,9 @@ def test_data_validation_error_messages(self, data_manager):
             # Should contain user-friendly keywords
             message_lower = message.lower()
             for keyword in expected_keywords:
-                assert keyword in message_lower, f"Error message should contain '{keyword}': {message}"
+                assert keyword in message_lower, (
+                    f"Error message should contain '{keyword}': {message}"
+                )
 
         print("✅ Data validation error messages test passed")
 
@@ -107,10 +109,10 @@ def test_empty_data_handling(self, calculator):
 
         for empty_df in empty_scenarios:
             steps = ["Step1", "Step2"]
-            
+
             # Should handle empty data gracefully
             results = calculator.calculate_funnel_metrics(empty_df, steps)
-            
+
             # Should return valid structure with zero counts
             assert results.steps == [], "Should return empty steps for empty data"
             assert results.users_count == [], "Should return empty users_count for empty data"
@@ -120,13 +122,15 @@ def test_empty_data_handling(self, calculator):
     def test_invalid_funnel_steps_handling(self, calculator):
         """Test handling of invalid funnel step configurations."""
         # Create valid data
-        valid_data = pd.DataFrame({
-            "user_id": ["user1", "user2", "user3"],
-            "event_name": ["EventA", "EventB", "EventC"],
-            "timestamp": [datetime.now() + timedelta(minutes=i) for i in range(3)],
-            "event_properties": ["{}"] * 3,
-            "user_properties": ["{}"] * 3,
-        })
+        valid_data = pd.DataFrame(
+            {
+                "user_id": ["user1", "user2", "user3"],
+                "event_name": ["EventA", "EventB", "EventC"],
+                "timestamp": [datetime.now() + timedelta(minutes=i) for i in range(3)],
+                "event_properties": ["{}"] * 3,
+                "user_properties": ["{}"] * 3,
+            }
+        )
 
         invalid_step_scenarios = [
             [],  # Empty steps
@@ -137,10 +141,12 @@ def test_invalid_funnel_steps_handling(self, calculator):
         for steps in invalid_step_scenarios:
             # Should handle invalid steps gracefully
             results = calculator.calculate_funnel_metrics(valid_data, steps)
-            
+
             # Should return appropriate structure
             if not steps:
-                assert results.steps == [], f"Should return empty steps for invalid configuration: {steps}"
+                assert results.steps == [], (
+                    f"Should return empty steps for invalid configuration: {steps}"
+                )
             else:
                 # Should filter out non-existent events
                 assert isinstance(results.steps, list), f"Should return list for steps: {steps}"
@@ -155,18 +161,20 @@ class TestRecoveryMechanisms:
     def test_partial_data_recovery(self):
         """Test recovery from partial data corruption."""
         # Create partially corrupted data
-        partial_data = pd.DataFrame({
-            "user_id": ["user1", None, "user3", ""],  # Some missing user IDs
-            "event_name": ["Event1", "Event2", None, "Event4"],  # Some missing events
-            "timestamp": [
-                datetime.now(),
-                datetime.now() + timedelta(minutes=1),
-                None,  # Missing timestamp
-                datetime.now() + timedelta(minutes=3),
-            ],
-            "event_properties": ["{}"] * 4,
-            "user_properties": ["{}"] * 4,
-        })
+        partial_data = pd.DataFrame(
+            {
+                "user_id": ["user1", None, "user3", ""],  # Some missing user IDs
+                "event_name": ["Event1", "Event2", None, "Event4"],  # Some missing events
+                "timestamp": [
+                    datetime.now(),
+                    datetime.now() + timedelta(minutes=1),
+                    None,  # Missing timestamp
+                    datetime.now() + timedelta(minutes=3),
+                ],
+                "event_properties": ["{}"] * 4,
+                "user_properties": ["{}"] * 4,
+            }
+        )
 
         config = FunnelConfig()
         calculator = FunnelCalculator(config)
@@ -174,7 +182,7 @@ def test_partial_data_recovery(self):
 
         # Should recover and process valid data
         results = calculator.calculate_funnel_metrics(partial_data, steps)
-        
+
         # Should have some results from valid data
         assert isinstance(results.steps, list), "Should return valid results from partial data"
 
@@ -188,7 +196,7 @@ class TestUserFriendlyErrorMessages:
     def test_error_message_clarity(self):
         """Test that error messages are clear and actionable."""
         data_manager = DataSourceManager()
-        
+
         # Test scenarios with expected user-friendly messages
         error_scenarios = [
             (
@@ -200,17 +208,19 @@ def test_error_message_clarity(self):
 
         for test_data, should_contain, should_not_contain in error_scenarios:
             is_valid, message = data_manager.validate_event_data(test_data)
-            
+
             assert not is_valid, "Should be invalid"
-            
+
             message_lower = message.lower()
-            
+
             # Should contain helpful keywords
             for keyword in should_contain:
                 assert keyword in message_lower, f"Message should contain '{keyword}': {message}"
-            
+
             # Should not contain technical jargon
             for jargon in should_not_contain:
-                assert jargon not in message_lower, f"Message should not contain '{jargon}': {message}"
+                assert jargon not in message_lower, (
+                    f"Message should not contain '{jargon}': {message}"
+                )
 
-        print("✅ Error message clarity test passed") 
\ No newline at end of file
+        print("✅ Error message clarity test passed")
diff --git a/tests/test_file_upload_comprehensive.py b/tests/test_file_upload_comprehensive.py
index 76ac495..e8ab2f5 100644
--- a/tests/test_file_upload_comprehensive.py
+++ b/tests/test_file_upload_comprehensive.py
@@ -14,12 +14,9 @@
 - @pytest.mark.security - Security validation tests
 """
 
-import io
 import json
-import tempfile
 import time
 from datetime import datetime, timedelta
-from pathlib import Path
 from unittest.mock import Mock, patch
 
 import pandas as pd
@@ -48,33 +45,41 @@ def valid_csv_content(self):
     @pytest.fixture
     def valid_parquet_data(self):
         """Valid DataFrame for Parquet testing."""
-        return pd.DataFrame({
-            "user_id": ["user_001", "user_002", "user_003"],
-            "event_name": ["Sign Up", "Login", "Purchase"],
-            "timestamp": [
-                datetime(2024, 1, 1, 10, 0, 0),
-                datetime(2024, 1, 1, 11, 0, 0),
-                datetime(2024, 1, 1, 12, 0, 0),
-            ],
-            "event_properties": ["{}", "{}", "{}"],
-            "user_properties": ["{}", "{}", "{}"],
-        })
+        return pd.DataFrame(
+            {
+                "user_id": ["user_001", "user_002", "user_003"],
+                "event_name": ["Sign Up", "Login", "Purchase"],
+                "timestamp": [
+                    datetime(2024, 1, 1, 10, 0, 0),
+                    datetime(2024, 1, 1, 11, 0, 0),
+                    datetime(2024, 1, 1, 12, 0, 0),
+                ],
+                "event_properties": ["{}", "{}", "{}"],
+                "user_properties": ["{}", "{}", "{}"],
+            }
+        )
 
     def test_csv_format_validation_success(self, data_manager):
         """Test successful CSV file validation and loading."""
         # Create mock uploaded file
         mock_file = Mock()
         mock_file.name = "test_events.csv"
-        
+
         with patch("pandas.read_csv") as mock_read_csv:
             # Setup pandas to return valid DataFrame
-            expected_df = pd.DataFrame({
-                "user_id": ["user_001", "user_002", "user_003"],
-                "event_name": ["Sign Up", "Login", "Purchase"],
-                "timestamp": ["2024-01-01 10:00:00", "2024-01-01 11:00:00", "2024-01-01 12:00:00"],
-                "event_properties": ["{}", "{}", "{}"],
-                "user_properties": ["{}", "{}", "{}"],
-            })
+            expected_df = pd.DataFrame(
+                {
+                    "user_id": ["user_001", "user_002", "user_003"],
+                    "event_name": ["Sign Up", "Login", "Purchase"],
+                    "timestamp": [
+                        "2024-01-01 10:00:00",
+                        "2024-01-01 11:00:00",
+                        "2024-01-01 12:00:00",
+                    ],
+                    "event_properties": ["{}", "{}", "{}"],
+                    "user_properties": ["{}", "{}", "{}"],
+                }
+            )
             mock_read_csv.return_value = expected_df
 
             # Load file
@@ -94,7 +99,7 @@ def test_parquet_format_validation_success(self, data_manager, valid_parquet_dat
         # Create mock uploaded file
         mock_file = Mock()
         mock_file.name = "test_events.parquet"
-        
+
         with patch("pandas.read_parquet") as mock_read_parquet:
             mock_read_parquet.return_value = valid_parquet_data
 
@@ -127,7 +132,7 @@ def test_unsupported_file_format_rejection(self, data_manager):
     def test_missing_required_columns_validation(self, data_manager):
         """Test validation of missing required columns."""
         required_columns = ["user_id", "event_name", "timestamp"]
-        
+
         for missing_col in required_columns:
             # Create data missing one required column
             test_data = {
@@ -138,14 +143,16 @@ def test_missing_required_columns_validation(self, data_manager):
                 "user_properties": ["{}", "{}"],
             }
             del test_data[missing_col]
-            
+
             invalid_df = pd.DataFrame(test_data)
-            
+
             # Test validation
             is_valid, message = data_manager.validate_event_data(invalid_df)
-            
+
             assert not is_valid, f"Should be invalid when missing {missing_col}"
-            assert missing_col.lower() in message.lower(), f"Error message should mention {missing_col}"
+            assert missing_col.lower() in message.lower(), (
+                f"Error message should mention {missing_col}"
+            )
 
         print("✅ Missing required columns validation test passed")
 
@@ -161,19 +168,23 @@ def test_invalid_timestamp_format_handling(self, data_manager):
         ]
 
         for timestamps in invalid_timestamp_scenarios:
-            df = pd.DataFrame({
-                "user_id": ["user1", "user2"],
-                "event_name": ["Event1", "Event2"],
-                "timestamp": timestamps,
-                "event_properties": ["{}", "{}"],
-                "user_properties": ["{}", "{}"],
-            })
+            df = pd.DataFrame(
+                {
+                    "user_id": ["user1", "user2"],
+                    "event_name": ["Event1", "Event2"],
+                    "timestamp": timestamps,
+                    "event_properties": ["{}", "{}"],
+                    "user_properties": ["{}", "{}"],
+                }
+            )
 
             is_valid, message = data_manager.validate_event_data(df)
 
             # Should either be invalid or successfully convert timestamps
             if not is_valid:
-                assert "timestamp" in message.lower(), f"Should mention timestamp issues for {timestamps}"
+                assert "timestamp" in message.lower(), (
+                    f"Should mention timestamp issues for {timestamps}"
+                )
 
         print("✅ Invalid timestamp format handling test passed")
 
@@ -224,7 +235,9 @@ def test_corrupted_parquet_file_handling(self, data_manager):
                 result_df = data_manager.load_from_file(mock_file)
 
                 # Should handle gracefully
-                assert isinstance(result_df, pd.DataFrame), f"Should return DataFrame for error: {error}"
+                assert isinstance(result_df, pd.DataFrame), (
+                    f"Should return DataFrame for error: {error}"
+                )
                 assert len(result_df) == 0, f"Should return empty DataFrame for error: {error}"
 
         print("✅ Corrupted Parquet file handling test passed")
@@ -233,22 +246,24 @@ def test_malformed_json_properties_handling(self, data_manager):
         """Test handling of malformed JSON in event/user properties."""
         malformed_json_scenarios = [
             '{"valid": "json"}',  # Valid JSON
-            "{invalid json}",     # Invalid JSON syntax
-            "not json at all",    # Not JSON
-            '{"incomplete": }',   # Incomplete JSON
-            "",                   # Empty string
-            None,                 # Null value
+            "{invalid json}",  # Invalid JSON syntax
+            "not json at all",  # Not JSON
+            '{"incomplete": }',  # Incomplete JSON
+            "",  # Empty string
+            None,  # Null value
             '{"nested": {"deep": {"very": "deep"}}}',  # Very nested JSON
         ]
 
         # Create DataFrame with malformed JSON
-        df = pd.DataFrame({
-            "user_id": [f"user_{i}" for i in range(len(malformed_json_scenarios))],
-            "event_name": ["Event"] * len(malformed_json_scenarios),
-            "timestamp": [datetime.now()] * len(malformed_json_scenarios),
-            "event_properties": malformed_json_scenarios,
-            "user_properties": malformed_json_scenarios,
-        })
+        df = pd.DataFrame(
+            {
+                "user_id": [f"user_{i}" for i in range(len(malformed_json_scenarios))],
+                "event_name": ["Event"] * len(malformed_json_scenarios),
+                "timestamp": [datetime.now()] * len(malformed_json_scenarios),
+                "event_properties": malformed_json_scenarios,
+                "user_properties": malformed_json_scenarios,
+            }
+        )
 
         # Should validate successfully (malformed JSON doesn't fail validation)
         is_valid, message = data_manager.validate_event_data(df)
@@ -268,7 +283,7 @@ def test_encoding_issues_handling(self, data_manager):
         with patch("pandas.read_csv") as mock_read_csv:
             # Simulate encoding errors
             encoding_errors = [
-                UnicodeDecodeError('utf-8', b'', 0, 1, 'invalid start byte'),
+                UnicodeDecodeError("utf-8", b"", 0, 1, "invalid start byte"),
                 UnicodeError("Unicode error"),
                 Exception("Encoding detection failed"),
             ]
@@ -279,8 +294,12 @@ def test_encoding_issues_handling(self, data_manager):
                 result_df = data_manager.load_from_file(mock_file)
 
                 # Should handle encoding errors gracefully
-                assert isinstance(result_df, pd.DataFrame), f"Should return DataFrame for encoding error: {error}"
-                assert len(result_df) == 0, f"Should return empty DataFrame for encoding error: {error}"
+                assert isinstance(result_df, pd.DataFrame), (
+                    f"Should return DataFrame for encoding error: {error}"
+                )
+                assert len(result_df) == 0, (
+                    f"Should return empty DataFrame for encoding error: {error}"
+                )
 
         print("✅ Encoding issues handling test passed")
 
@@ -303,13 +322,17 @@ def test_large_csv_file_performance(self, data_manager):
 
         with patch("pandas.read_csv") as mock_read_csv:
             # Create large DataFrame
-            large_df = pd.DataFrame({
-                "user_id": [f"user_{i:06d}" for i in range(n_rows)],
-                "event_name": [f"Event_{i % 10}" for i in range(n_rows)],
-                "timestamp": [datetime(2024, 1, 1) + timedelta(minutes=i) for i in range(n_rows)],
-                "event_properties": ["{}"] * n_rows,
-                "user_properties": ["{}"] * n_rows,
-            })
+            large_df = pd.DataFrame(
+                {
+                    "user_id": [f"user_{i:06d}" for i in range(n_rows)],
+                    "event_name": [f"Event_{i % 10}" for i in range(n_rows)],
+                    "timestamp": [
+                        datetime(2024, 1, 1) + timedelta(minutes=i) for i in range(n_rows)
+                    ],
+                    "event_properties": ["{}"] * n_rows,
+                    "user_properties": ["{}"] * n_rows,
+                }
+            )
             mock_read_csv.return_value = large_df
 
             # Measure loading time
@@ -320,46 +343,62 @@ def test_large_csv_file_performance(self, data_manager):
             execution_time = end_time - start_time
 
             # Should load in reasonable time (less than 10 seconds)
-            assert execution_time < 10.0, f"Large file loading took too long: {execution_time:.2f}s"
-            
+            assert execution_time < 10.0, (
+                f"Large file loading took too long: {execution_time:.2f}s"
+            )
+
             # Should load successfully
             assert isinstance(result_df, pd.DataFrame), "Should return DataFrame"
             assert len(result_df) == n_rows, f"Should load all {n_rows} rows"
 
-        print(f"✅ Large CSV file performance test passed ({execution_time:.2f}s for {n_rows} rows)")
+        print(
+            f"✅ Large CSV file performance test passed ({execution_time:.2f}s for {n_rows} rows)"
+        )
 
     def test_memory_usage_monitoring(self, data_manager):
         """Test memory usage with large datasets."""
         import tracemalloc
-        
+
         # Start memory tracking
         tracemalloc.start()
-        
+
         n_rows = 100000
         mock_file = Mock()
         mock_file.name = "memory_test.csv"
 
         with patch("pandas.read_csv") as mock_read_csv:
             # Create memory-intensive DataFrame
-            large_df = pd.DataFrame({
-                "user_id": [f"user_{i:06d}" for i in range(n_rows)],
-                "event_name": [f"Event_{i % 20}" for i in range(n_rows)],
-                "timestamp": [datetime(2024, 1, 1) + timedelta(minutes=i) for i in range(n_rows)],
-                "event_properties": [
-                    json.dumps({"category": f"cat_{i % 10}", "value": i, "metadata": {"nested": True}})
-                    for i in range(n_rows)
-                ],
-                "user_properties": [
-                    json.dumps({"segment": f"seg_{i % 5}", "score": i % 100, "profile": {"active": True}})
-                    for i in range(n_rows)
-                ],
-            })
+            large_df = pd.DataFrame(
+                {
+                    "user_id": [f"user_{i:06d}" for i in range(n_rows)],
+                    "event_name": [f"Event_{i % 20}" for i in range(n_rows)],
+                    "timestamp": [
+                        datetime(2024, 1, 1) + timedelta(minutes=i) for i in range(n_rows)
+                    ],
+                    "event_properties": [
+                        json.dumps(
+                            {"category": f"cat_{i % 10}", "value": i, "metadata": {"nested": True}}
+                        )
+                        for i in range(n_rows)
+                    ],
+                    "user_properties": [
+                        json.dumps(
+                            {
+                                "segment": f"seg_{i % 5}",
+                                "score": i % 100,
+                                "profile": {"active": True},
+                            }
+                        )
+                        for i in range(n_rows)
+                    ],
+                }
+            )
             mock_read_csv.return_value = large_df
 
             # Load and validate data
             result_df = data_manager.load_from_file(mock_file)
             is_valid, _ = data_manager.validate_event_data(result_df)
-            
+
             # Get memory usage
             current, peak = tracemalloc.get_traced_memory()
             tracemalloc.stop()
@@ -380,7 +419,7 @@ def test_chunked_processing_simulation(self, data_manager):
         chunk_size = 10000
         total_rows = 50000
         chunks_processed = 0
-        
+
         mock_file = Mock()
         mock_file.name = "chunked_test.csv"
 
@@ -388,31 +427,39 @@ def test_chunked_processing_simulation(self, data_manager):
         for chunk_start in range(0, total_rows, chunk_size):
             chunk_end = min(chunk_start + chunk_size, total_rows)
             chunk_rows = chunk_end - chunk_start
-            
+
             with patch("pandas.read_csv") as mock_read_csv:
                 # Create chunk DataFrame
-                chunk_df = pd.DataFrame({
-                    "user_id": [f"user_{i:06d}" for i in range(chunk_start, chunk_end)],
-                    "event_name": [f"Event_{i % 5}" for i in range(chunk_rows)],
-                    "timestamp": [datetime(2024, 1, 1) + timedelta(minutes=i) for i in range(chunk_rows)],
-                    "event_properties": ["{}"] * chunk_rows,
-                    "user_properties": ["{}"] * chunk_rows,
-                })
+                chunk_df = pd.DataFrame(
+                    {
+                        "user_id": [f"user_{i:06d}" for i in range(chunk_start, chunk_end)],
+                        "event_name": [f"Event_{i % 5}" for i in range(chunk_rows)],
+                        "timestamp": [
+                            datetime(2024, 1, 1) + timedelta(minutes=i) for i in range(chunk_rows)
+                        ],
+                        "event_properties": ["{}"] * chunk_rows,
+                        "user_properties": ["{}"] * chunk_rows,
+                    }
+                )
                 mock_read_csv.return_value = chunk_df
 
                 # Process chunk
                 result_df = data_manager.load_from_file(mock_file)
                 is_valid, _ = data_manager.validate_event_data(result_df)
-                
+
                 assert is_valid, f"Chunk {chunks_processed} should be valid"
-                assert len(result_df) == chunk_rows, f"Chunk {chunks_processed} should have {chunk_rows} rows"
-                
+                assert len(result_df) == chunk_rows, (
+                    f"Chunk {chunks_processed} should have {chunk_rows} rows"
+                )
+
                 chunks_processed += 1
 
         expected_chunks = (total_rows + chunk_size - 1) // chunk_size
         assert chunks_processed == expected_chunks, f"Should process {expected_chunks} chunks"
 
-        print(f"✅ Chunked processing simulation test passed ({chunks_processed} chunks of {chunk_size} rows)")
+        print(
+            f"✅ Chunked processing simulation test passed ({chunks_processed} chunks of {chunk_size} rows)"
+        )
 
     def test_concurrent_file_upload_simulation(self, data_manager):
         """Test sequential simulation of concurrent file uploads (avoiding threading issues in tests)."""
@@ -428,28 +475,36 @@ def test_concurrent_file_upload_simulation(self, data_manager):
 
             with patch("pandas.read_csv") as mock_read_csv:
                 # Create test DataFrame
-                test_df = pd.DataFrame({
-                    "user_id": [f"worker_{worker_id}_user_{i}" for i in range(500)],  # Smaller size
-                    "event_name": [f"Event_{i % 3}" for i in range(500)],
-                    "timestamp": [datetime.now() + timedelta(minutes=i) for i in range(500)],
-                    "event_properties": ["{}"] * 500,
-                    "user_properties": ["{}"] * 500,
-                })
+                test_df = pd.DataFrame(
+                    {
+                        "user_id": [
+                            f"worker_{worker_id}_user_{i}" for i in range(500)
+                        ],  # Smaller size
+                        "event_name": [f"Event_{i % 3}" for i in range(500)],
+                        "timestamp": [datetime.now() + timedelta(minutes=i) for i in range(500)],
+                        "event_properties": ["{}"] * 500,
+                        "user_properties": ["{}"] * 500,
+                    }
+                )
                 mock_read_csv.return_value = test_df
 
                 start_time = time.time()
                 result_df = data_manager.load_from_file(mock_file)
                 end_time = time.time()
 
-                results.append({
-                    "worker_id": worker_id,
-                    "duration": end_time - start_time,
-                    "rows_loaded": len(result_df),
-                    "success": len(result_df) == 500,
-                })
+                results.append(
+                    {
+                        "worker_id": worker_id,
+                        "duration": end_time - start_time,
+                        "rows_loaded": len(result_df),
+                        "success": len(result_df) == 500,
+                    }
+                )
 
         # Validate results
-        assert len(results) == num_simulated_uploads, f"All {num_simulated_uploads} simulated uploads should complete"
+        assert len(results) == num_simulated_uploads, (
+            f"All {num_simulated_uploads} simulated uploads should complete"
+        )
 
         # Check performance
         avg_duration = sum(r["duration"] for r in results) / len(results)
@@ -457,9 +512,13 @@ def test_concurrent_file_upload_simulation(self, data_manager):
 
         # Check success
         successful_uploads = sum(1 for r in results if r["success"])
-        assert successful_uploads == num_simulated_uploads, f"All simulated uploads should succeed, got {successful_uploads}/{num_simulated_uploads}"
+        assert successful_uploads == num_simulated_uploads, (
+            f"All simulated uploads should succeed, got {successful_uploads}/{num_simulated_uploads}"
+        )
 
-        print(f"✅ Concurrent file upload simulation test passed ({num_simulated_uploads} simulated uploads, avg {avg_duration:.2f}s)")
+        print(
+            f"✅ Concurrent file upload simulation test passed ({num_simulated_uploads} simulated uploads, avg {avg_duration:.2f}s)"
+        )
 
 
 @pytest.mark.file_upload
@@ -495,7 +554,9 @@ def test_filename_validation_security(self, data_manager):
             result_df = data_manager.load_from_file(mock_file)
 
             # Should return empty DataFrame for security (since we don't actually read the file)
-            assert isinstance(result_df, pd.DataFrame), f"Should return DataFrame for filename: {filename}"
+            assert isinstance(result_df, pd.DataFrame), (
+                f"Should return DataFrame for filename: {filename}"
+            )
             # Note: Actual file reading is mocked, so we get empty DataFrame
 
         print("✅ Filename validation security test passed")
@@ -512,7 +573,9 @@ def test_file_size_limits_simulation(self, data_manager):
             result_df = data_manager.load_from_file(mock_file)
 
             # Should handle memory errors gracefully
-            assert isinstance(result_df, pd.DataFrame), "Should return DataFrame even for memory errors"
+            assert isinstance(result_df, pd.DataFrame), (
+                "Should return DataFrame even for memory errors"
+            )
             assert len(result_df) == 0, "Should return empty DataFrame for oversized files"
 
         print("✅ File size limits simulation test passed")
@@ -522,38 +585,44 @@ def test_content_validation_security(self, data_manager):
         # Test with potentially malicious content
         malicious_content_scenarios = [
             # SQL injection attempts in data
-            pd.DataFrame({
-                "user_id": ["'; DROP TABLE users; --", "user2"],
-                "event_name": ["SELECT * FROM events", "Event2"],
-                "timestamp": [datetime.now(), datetime.now()],
-                "event_properties": ["{}"] * 2,
-                "user_properties": ["{}"] * 2,
-            }),
-            
+            pd.DataFrame(
+                {
+                    "user_id": ["'; DROP TABLE users; --", "user2"],
+                    "event_name": ["SELECT * FROM events", "Event2"],
+                    "timestamp": [datetime.now(), datetime.now()],
+                    "event_properties": ["{}"] * 2,
+                    "user_properties": ["{}"] * 2,
+                }
+            ),
             # XSS attempts in data
-            pd.DataFrame({
-                "user_id": ["<script>alert('xss')</script>", "user2"],
-                "event_name": ["<img src=x onerror=alert(1)>", "Event2"],
-                "timestamp": [datetime.now(), datetime.now()],
-                "event_properties": ["{}"] * 2,
-                "user_properties": ["{}"] * 2,
-            }),
-            
+            pd.DataFrame(
+                {
+                    "user_id": ["<script>alert('xss')</script>", "user2"],
+                    "event_name": ["<img src=x onerror=alert(1)>", "Event2"],
+                    "timestamp": [datetime.now(), datetime.now()],
+                    "event_properties": ["{}"] * 2,
+                    "user_properties": ["{}"] * 2,
+                }
+            ),
             # Extremely long strings (potential buffer overflow)
-            pd.DataFrame({
-                "user_id": ["A" * 10000, "user2"],
-                "event_name": ["B" * 10000, "Event2"],
-                "timestamp": [datetime.now(), datetime.now()],
-                "event_properties": ["{}"] * 2,
-                "user_properties": ["{}"] * 2,
-            }),
+            pd.DataFrame(
+                {
+                    "user_id": ["A" * 10000, "user2"],
+                    "event_name": ["B" * 10000, "Event2"],
+                    "timestamp": [datetime.now(), datetime.now()],
+                    "event_properties": ["{}"] * 2,
+                    "user_properties": ["{}"] * 2,
+                }
+            ),
         ]
 
         for i, malicious_df in enumerate(malicious_content_scenarios):
             # Should validate without crashing
             is_valid, message = data_manager.validate_event_data(malicious_df)
-            
+
             # Should be valid structurally (content filtering is not part of validation)
-            assert is_valid or "timestamp" in message.lower(), f"Scenario {i} should validate or have timestamp issue"
+            assert is_valid or "timestamp" in message.lower(), (
+                f"Scenario {i} should validate or have timestamp issue"
+            )
 
-        print("✅ Content validation security test passed") 
\ No newline at end of file
+        print("✅ Content validation security test passed")
diff --git a/tests/test_visualization_pipeline_comprehensive.py b/tests/test_visualization_pipeline_comprehensive.py
index 1eb8ca5..8d7aeb1 100644
--- a/tests/test_visualization_pipeline_comprehensive.py
+++ b/tests/test_visualization_pipeline_comprehensive.py
@@ -10,19 +10,12 @@
 """
 
 import json
-import time
-from datetime import datetime, timedelta
-from unittest.mock import Mock, patch
 
-import pandas as pd
 import plotly.graph_objects as go
 import pytest
 
 from models import (
-    CohortData,
     FunnelResults,
-    PathAnalysisData,
-    TimeToConvertStats,
 )
 from ui.visualization import FunnelVisualizer
 from ui.visualization.layout import LayoutConfig
@@ -59,21 +52,23 @@ def test_funnel_chart_rendering_pipeline(self, visualizer, sample_funnel_results
         """Test complete funnel chart rendering pipeline."""
         # Test basic funnel chart
         chart = visualizer.create_enhanced_funnel_chart(sample_funnel_results)
-        
+
         # Validate chart object
         assert isinstance(chart, go.Figure), "Should return Plotly Figure object"
         assert chart.data is not None, "Chart should have data"
         assert chart.layout is not None, "Chart should have layout"
-        
+
         # Validate data traces
         assert len(chart.data) > 0, "Chart should have at least one trace"
-        
+
         # Validate layout properties
         assert chart.layout.title is not None, "Chart should have title"
         assert chart.layout.height is not None, "Chart should have height"
-        
+
         # Test with segments
-        segmented_chart = visualizer.create_enhanced_funnel_chart(sample_funnel_results, show_segments=True)
+        segmented_chart = visualizer.create_enhanced_funnel_chart(
+            sample_funnel_results, show_segments=True
+        )
         assert isinstance(segmented_chart, go.Figure), "Segmented chart should be valid Figure"
 
         print("✅ Funnel chart rendering pipeline test passed")
@@ -81,20 +76,20 @@ def test_funnel_chart_rendering_pipeline(self, visualizer, sample_funnel_results
     def test_sankey_diagram_rendering_pipeline(self, visualizer, sample_funnel_results):
         """Test Sankey diagram rendering pipeline."""
         chart = visualizer.create_enhanced_conversion_flow_sankey(sample_funnel_results)
-        
+
         # Validate Sankey structure
         assert isinstance(chart, go.Figure), "Should return Plotly Figure object"
-        
+
         # Find Sankey trace
         sankey_trace = None
         for trace in chart.data:
-            if hasattr(trace, 'type') and trace.type == 'sankey':
+            if hasattr(trace, "type") and trace.type == "sankey":
                 sankey_trace = trace
                 break
-        
+
         assert sankey_trace is not None, "Should contain Sankey trace"
-        assert hasattr(sankey_trace, 'node'), "Sankey should have nodes"
-        assert hasattr(sankey_trace, 'link'), "Sankey should have links"
+        assert hasattr(sankey_trace, "node"), "Sankey should have nodes"
+        assert hasattr(sankey_trace, "link"), "Sankey should have links"
 
         print("✅ Sankey diagram rendering pipeline test passed")
 
@@ -113,16 +108,18 @@ def test_responsive_height_calculation(self, visualizer):
         """Test responsive height calculation algorithm."""
         # Test LayoutConfig responsive height calculation
         test_scenarios = [
-            (400, 3, 400),    # Small funnel, minimum height
-            (400, 10, 580),   # Medium funnel, scaled height
-            (400, 50, 800),   # Large funnel, capped at maximum
+            (400, 3, 400),  # Small funnel, minimum height
+            (400, 10, 580),  # Medium funnel, scaled height
+            (400, 50, 800),  # Large funnel, capped at maximum
         ]
-        
+
         for base_height, data_items, expected_range in test_scenarios:
             calculated_height = LayoutConfig.get_responsive_height(base_height, data_items)
-            
+
             # Should be within universal height standards
-            assert 350 <= calculated_height <= 800, f"Height {calculated_height} outside standards for {data_items} items"
+            assert 350 <= calculated_height <= 800, (
+                f"Height {calculated_height} outside standards for {data_items} items"
+            )
 
         print("✅ Responsive height calculation test passed")
 
@@ -136,12 +133,12 @@ def test_mobile_compatibility(self, visualizer):
             drop_offs=[0, 20, 20],
             drop_off_rates=[0.0, 20.0, 25.0],
         )
-        
+
         chart = visualizer.create_enhanced_funnel_chart(mobile_results)
-        
+
         # Validate mobile-friendly properties
         assert chart.layout.height <= 600, "Mobile charts should have reasonable height"
-        
+
         # Check margin configuration for mobile
         margins = chart.layout.margin
         assert margins.l <= 80, "Left margin should be mobile-friendly"
@@ -174,15 +171,15 @@ def test_colorblind_friendly_palettes(self, visualizer):
             "Segment A": [600, 500, 400],
             "Segment B": [400, 300, 200],
         }
-        
+
         chart = visualizer.create_enhanced_funnel_chart(segmented_results, show_segments=True)
-        
+
         # Validate colorblind-friendly colors are used
         colors_used = []
         for trace in chart.data:
-            if hasattr(trace, 'marker') and trace.marker and hasattr(trace.marker, 'color'):
+            if hasattr(trace, "marker") and trace.marker and hasattr(trace.marker, "color"):
                 colors_used.append(trace.marker.color)
-        
+
         # Should use distinct, accessible colors
         assert len(set(colors_used)) >= 2, "Should use multiple distinct colors"
 
@@ -200,45 +197,49 @@ def visualizer(self):
 
     def test_plotly_figure_structure(self, visualizer):
         """Test Plotly Figure structure compliance."""
-        chart = visualizer.create_enhanced_funnel_chart(FunnelResults(
-            steps=["Step 1", "Step 2"],
-            users_count=[100, 80],
-            conversion_rates=[100.0, 80.0],
-            drop_offs=[0, 20],
-            drop_off_rates=[0.0, 20.0],
-        ))
-        
+        chart = visualizer.create_enhanced_funnel_chart(
+            FunnelResults(
+                steps=["Step 1", "Step 2"],
+                users_count=[100, 80],
+                conversion_rates=[100.0, 80.0],
+                drop_offs=[0, 20],
+                drop_off_rates=[0.0, 20.0],
+            )
+        )
+
         # Validate Figure structure
         assert isinstance(chart, go.Figure), "Should be Plotly Figure object"
-        assert hasattr(chart, 'data'), "Should have data attribute"
-        assert hasattr(chart, 'layout'), "Should have layout attribute"
-        
+        assert hasattr(chart, "data"), "Should have data attribute"
+        assert hasattr(chart, "layout"), "Should have layout attribute"
+
         # Validate data traces
         assert isinstance(chart.data, tuple), "Data should be tuple of traces"
         for trace in chart.data:
-            assert hasattr(trace, 'type'), "Each trace should have type"
+            assert hasattr(trace, "type"), "Each trace should have type"
 
         print("✅ Plotly Figure structure test passed")
 
     def test_json_serialization_compatibility(self, visualizer):
         """Test JSON serialization compatibility."""
-        chart = visualizer.create_enhanced_funnel_chart(FunnelResults(
-            steps=["Step 1", "Step 2"],
-            users_count=[100, 80],
-            conversion_rates=[100.0, 80.0],
-            drop_offs=[0, 20],
-            drop_off_rates=[0.0, 20.0],
-        ))
-        
+        chart = visualizer.create_enhanced_funnel_chart(
+            FunnelResults(
+                steps=["Step 1", "Step 2"],
+                users_count=[100, 80],
+                conversion_rates=[100.0, 80.0],
+                drop_offs=[0, 20],
+                drop_off_rates=[0.0, 20.0],
+            )
+        )
+
         # Should be JSON serializable
         try:
             chart_json = chart.to_json()
             assert isinstance(chart_json, str), "Should serialize to JSON string"
-            
+
             # Should be valid JSON
             parsed_json = json.loads(chart_json)
             assert isinstance(parsed_json, dict), "Should parse to dictionary"
-            
+
         except Exception as e:
             assert False, f"Chart should be JSON serializable: {str(e)}"
 
@@ -263,10 +264,10 @@ def test_empty_data_visualization(self, visualizer):
             drop_offs=[],
             drop_off_rates=[],
         )
-        
+
         # Should handle empty data gracefully
         chart = visualizer.create_enhanced_funnel_chart(empty_results)
-        
+
         assert isinstance(chart, go.Figure), "Should return valid Figure for empty data"
 
         print("✅ Empty data visualization test passed")
@@ -283,4 +284,4 @@ def test_invalid_data_visualization(self, visualizer):
             # Expected behavior - should raise clear exception
             assert str(e), "Exception should have descriptive message"
 
-        print("✅ Invalid data visualization test passed") 
\ No newline at end of file
+        print("✅ Invalid data visualization test passed")