Improve the test script

tests/test_onboarding_flow.sh

@@ -12,6 +12,10 @@ TEST_EMAIL="onboarding.test.$(date +%s)@bakery.com"
 TEST_PASSWORD="TestPassword123!"
 TEST_NAME="Test Bakery Owner"
 REAL_CSV_FILE="bakery_sales_2023_2024.csv"
+WS_BASE="ws://localhost:8002/api/v1/ws"
+WS_TEST_DURATION=30  # seconds to listen for WebSocket messages
+WS_PID=""
+
 
 # Colors for output
 RED='\033[0;31m'
@@ -154,6 +158,591 @@ check_timezone_error() {
     return 1  # No timezone error
 }
 
+# Function to test WebSocket connection using websocat (if available) or Node.js
+test_websocket_connection() {
+    local tenant_id="$1"
+    local job_id="$2"
+    local duration="$3"
+    
+    log_step "4.2. Testing WebSocket connection for real-time training progress"
+    
+    echo "WebSocket URL: $WS_BASE/tenants/$tenant_id/training/jobs/$job_id/live"
+    echo "Test duration: ${duration}s"
+    echo ""
+    
+    # Check if websocat is available
+    if command -v websocat >/dev/null 2>&1; then
+        test_websocket_with_websocat "$tenant_id" "$job_id" "$duration"
+    elif command -v node >/dev/null 2>&1; then
+        test_websocket_with_nodejs_builtin "$tenant_id" "$job_id" "$duration"
+    else
+        test_websocket_with_curl "$tenant_id" "$job_id" "$duration"
+    fi
+}
+
+# Test WebSocket using websocat (recommended)
+test_websocket_with_websocat() {
+    local tenant_id="$1"
+    local job_id="$2"
+    local duration="$3"
+    
+    echo "Using websocat for WebSocket testing..."
+    
+    # Create a temporary file for WebSocket messages
+    local ws_log="/tmp/websocket_messages_$job_id.log"
+    
+    # Start WebSocket connection in background
+    (
+        echo "Connecting to WebSocket..."
+        timeout "${duration}s" websocat "$WS_BASE/tenants/$tenant_id/training/jobs/$job_id/live" \
+            --header "Authorization: Bearer $ACCESS_TOKEN" 2>&1 | \
+        while IFS= read -r line; do
+            echo "$(date '+%H:%M:%S') | $line" | tee -a "$ws_log"
+        done
+    ) &
+    
+    WS_PID=$!
+    
+    # Send periodic ping messages to keep connection alive
+    sleep 2
+    if kill -0 $WS_PID 2>/dev/null; then
+        echo "ping" | websocat "$WS_BASE/tenants/$tenant_id/training/jobs/$job_id/live" \
+            --header "Authorization: Bearer $ACCESS_TOKEN" >/dev/null 2>&1 &
+    fi
+    
+    # Wait for test duration
+    log_step "4.2.1. Listening for WebSocket messages (${duration}s)..."
+    wait_for_websocket_messages "$ws_log" "$duration"
+    
+    # Clean up
+    if kill -0 $WS_PID 2>/dev/null; then
+        kill $WS_PID 2>/dev/null
+        wait $WS_PID 2>/dev/null
+    fi
+}
+
+# Test WebSocket using Node.js
+test_websocket_with_nodejs() {
+    local tenant_id="$1"
+    local job_id="$2"
+    local duration="$3"
+    
+    echo "Using Node.js for WebSocket testing..."
+    
+    # Create Node.js WebSocket test script
+    local ws_test_script="/tmp/websocket_test_$job_id.js"
+    cat > "$ws_test_script" << 'EOF'
+const WebSocket = require('ws');
+
+const tenantId = process.argv[2];
+const jobId = process.argv[3];
+const duration = parseInt(process.argv[4]) * 1000;
+const accessToken = process.argv[5];
+const wsUrl = process.argv[6];
+
+console.log(`Connecting to: ${wsUrl}`);
+
+const ws = new WebSocket(wsUrl, {
+    headers: {
+        'Authorization': `Bearer ${accessToken}`
+    }
+});
+
+let messageCount = 0;
+let startTime = Date.now();
+
+ws.on('open', function() {
+    console.log('✅ WebSocket connected successfully');
+    
+    // Send periodic pings
+    const pingInterval = setInterval(() => {
+        if (ws.readyState === WebSocket.OPEN) {
+            ws.send('ping');
+        }
+    }, 5000);
+    
+    // Close after duration
+    setTimeout(() => {
+        clearInterval(pingInterval);
+        console.log(`\n📊 WebSocket test completed after ${duration/1000}s`);
+        console.log(`📨 Total messages received: ${messageCount}`);
+        if (messageCount > 0) {
+            console.log('✅ WebSocket communication successful');
+        } else {
+            console.log('⚠️  No training progress messages received');
+            console.log('   This may be normal if training completed quickly');
+        }
+        ws.close();
+        process.exit(0);
+    }, duration);
+});
+
+ws.on('message', function(data) {
+    messageCount++;
+    const timestamp = new Date().toLocaleTimeString();
+    
+    try {
+        const message = JSON.parse(data);
+        console.log(`\n[${timestamp}] 📨 Message ${messageCount}:`);
+        console.log(`   Type: ${message.type || 'unknown'}`);
+        console.log(`   Job ID: ${message.job_id || 'unknown'}`);
+        
+        if (message.data) {
+            if (message.data.progress !== undefined) {
+                console.log(`   Progress: ${message.data.progress}%`);
+            }
+            if (message.data.current_step) {
+                console.log(`   Step: ${message.data.current_step}`);
+            }
+            if (message.data.current_product) {
+                console.log(`   Product: ${message.data.current_product}`);
+            }
+            if (message.data.estimated_time_remaining_minutes) {
+                console.log(`   ETA: ${message.data.estimated_time_remaining_minutes} minutes`);
+            }
+        }
+        
+        // Special handling for completion messages
+        if (message.type === 'completed') {
+            console.log('🎉 Training completed!');
+        } else if (message.type === 'failed') {
+            console.log('❌ Training failed!');
+        }
+        
+    } catch (e) {
+        console.log(`[${timestamp}] Raw message: ${data}`);
+    }
+});
+
+ws.on('error', function(error) {
+    console.log('❌ WebSocket error:', error.message);
+});
+
+ws.on('close', function(code, reason) {
+    console.log(`\n🔌 WebSocket closed (code: ${code}, reason: ${reason || 'normal'})`);
+    process.exit(code === 1000 ? 0 : 1);
+});
+EOF
+
+    # Run Node.js WebSocket test
+    local ws_url="$WS_BASE/tenants/$tenant_id/training/jobs/$job_id/live"
+    node "$ws_test_script" "$tenant_id" "$job_id" "$duration" "$ACCESS_TOKEN" "$ws_url" &
+    WS_PID=$!
+    
+    # Wait for completion
+    wait $WS_PID
+    local exit_code=$?
+    
+    # Clean up
+    rm -f "$ws_test_script"
+    
+    if [ $exit_code -eq 0 ]; then
+        log_success "WebSocket test completed successfully"
+    else
+        log_warning "WebSocket test completed with issues"
+    fi
+}
+
+test_websocket_with_nodejs_builtin() {
+    local tenant_id="$1"
+    local job_id="$2"
+    local duration="$3"
+    
+    echo "Using Node.js with built-in modules for WebSocket testing..."
+    
+    # Create Node.js WebSocket test script using built-in modules only
+    local ws_test_script="/tmp/websocket_test_$job_id.js"
+    cat > "$ws_test_script" << 'EOF'
+// WebSocket test using only built-in Node.js modules
+const { WebSocket } = require('node:http');
+const https = require('https');
+const http = require('http');
+const crypto = require('crypto');
+
+const tenantId = process.argv[2];
+const jobId = process.argv[3];
+const duration = parseInt(process.argv[4]) * 1000;
+const accessToken = process.argv[5];
+const wsUrl = process.argv[6];
+
+console.log(`Connecting to: ${wsUrl}`);
+console.log(`Duration: ${duration/1000}s`);
+
+// Parse WebSocket URL
+const url = new URL(wsUrl);
+const isSecure = url.protocol === 'wss:';
+const port = url.port || (isSecure ? 443 : 80);
+
+// Create WebSocket key (required for WebSocket handshake)
+const key = crypto.randomBytes(16).toString('base64');
+
+// WebSocket handshake headers
+const headers = {
+    'Upgrade': 'websocket',
+    'Connection': 'Upgrade',
+    'Sec-WebSocket-Key': key,
+    'Sec-WebSocket-Version': '13',
+    'Authorization': `Bearer ${accessToken}`
+};
+
+const options = {
+    hostname: url.hostname,
+    port: port,
+    path: url.pathname,
+    method: 'GET',
+    headers: headers
+};
+
+console.log(`Attempting WebSocket handshake to ${url.hostname}:${port}${url.pathname}`);
+
+const client = isSecure ? https : http;
+let messageCount = 0;
+let startTime = Date.now();
+
+const req = client.request(options);
+
+req.on('upgrade', (res, socket, head) => {
+    console.log('✅ WebSocket handshake successful');
+    
+    let buffer = '';
+    
+    socket.on('data', (data) => {
+        buffer += data.toString();
+        
+        // Process complete WebSocket frames
+        while (buffer.length > 0) {
+            // Simple WebSocket frame parsing (for text frames)
+            if (buffer.length < 2) break;
+            
+            const firstByte = buffer.charCodeAt(0);
+            const secondByte = buffer.charCodeAt(1);
+            
+            const opcode = firstByte & 0x0F;
+            const masked = (secondByte & 0x80) === 0x80;
+            let payloadLength = secondByte & 0x7F;
+            
+            let offset = 2;
+            
+            if (payloadLength === 126) {
+                if (buffer.length < offset + 2) break;
+                payloadLength = (buffer.charCodeAt(offset) << 8) | buffer.charCodeAt(offset + 1);
+                offset += 2;
+            } else if (payloadLength === 127) {
+                if (buffer.length < offset + 8) break;
+                // For simplicity, assume payload length fits in 32 bits
+                payloadLength = (buffer.charCodeAt(offset + 4) << 24) | 
+                               (buffer.charCodeAt(offset + 5) << 16) | 
+                               (buffer.charCodeAt(offset + 6) << 8) | 
+                               buffer.charCodeAt(offset + 7);
+                offset += 8;
+            }
+            
+            if (buffer.length < offset + payloadLength) break;
+            
+            // Extract payload
+            let payload = buffer.slice(offset, offset + payloadLength);
+            buffer = buffer.slice(offset + payloadLength);
+            
+            if (opcode === 1) { // Text frame
+                messageCount++;
+                const timestamp = new Date().toLocaleTimeString();
+                
+                try {
+                    const message = JSON.parse(payload);
+                    console.log(`\n[${timestamp}] 📨 Message ${messageCount}:`);
+                    console.log(`   Type: ${message.type || 'unknown'}`);
+                    console.log(`   Job ID: ${message.job_id || 'unknown'}`);
+                    
+                    if (message.data) {
+                        if (message.data.progress !== undefined) {
+                            console.log(`   Progress: ${message.data.progress}%`);
+                        }
+                        if (message.data.current_step) {
+                            console.log(`   Step: ${message.data.current_step}`);
+                        }
+                    }
+                    
+                    if (message.type === 'completed') {
+                        console.log('🎉 Training completed!');
+                    } else if (message.type === 'failed') {
+                        console.log('❌ Training failed!');
+                    }
+                    
+                } catch (e) {
+                    console.log(`[${timestamp}] Raw message: ${payload}`);
+                }
+            } else if (opcode === 8) { // Close frame
+                console.log('🔌 WebSocket closed by server');
+                socket.end();
+                return;
+            }
+        }
+    });
+    
+    socket.on('end', () => {
+        console.log(`\n📊 WebSocket test completed`);
+        console.log(`📨 Total messages received: ${messageCount}`);
+        if (messageCount > 0) {
+            console.log('✅ WebSocket communication successful');
+        } else {
+            console.log('⚠️  No messages received during test period');
+        }
+        process.exit(0);
+    });
+    
+    socket.on('error', (error) => {
+        console.log('❌ WebSocket error:', error.message);
+        process.exit(1);
+    });
+    
+    // Send periodic pings to keep connection alive
+    const pingInterval = setInterval(() => {
+        if (socket.writable) {
+            // Send ping frame (opcode 9)
+            const pingFrame = Buffer.from([0x89, 0x00]);
+            socket.write(pingFrame);
+        }
+    }, 10000);
+    
+    // Close after duration
+    setTimeout(() => {
+        clearInterval(pingInterval);
+        console.log(`\n⏰ Test duration (${duration/1000}s) completed`);
+        console.log(`📨 Total messages received: ${messageCount}`);
+        
+        if (messageCount > 0) {
+            console.log('✅ WebSocket communication successful');
+        } else {
+            console.log('⚠️  No training progress messages received');
+            console.log('   This is normal if training completed before WebSocket connection');
+        }
+        
+        socket.end();
+        process.exit(0);
+    }, duration);
+});
+
+req.on('response', (res) => {
+    console.log(`❌ HTTP response instead of WebSocket upgrade: ${res.statusCode}`);
+    console.log('Response headers:', res.headers);
+    
+    let body = '';
+    res.on('data', chunk => body += chunk);
+    res.on('end', () => {
+        if (body) console.log('Response body:', body);
+    });
+    
+    process.exit(1);
+});
+
+req.on('error', (error) => {
+    console.log('❌ Connection error:', error.message);
+    process.exit(1);
+});
+
+req.end();
+EOF
+
+    # Run the Node.js WebSocket test
+    local ws_url="$WS_BASE/tenants/$tenant_id/training/jobs/$job_id/live"
+    echo "Starting WebSocket test..."
+    node "$ws_test_script" "$tenant_id" "$job_id" "$duration" "$ACCESS_TOKEN" "$ws_url"
+    local exit_code=$?
+    
+    # Clean up
+    rm -f "$ws_test_script"
+    
+    return $exit_code
+}
+
+# Fallback: Test WebSocket using curl (limited functionality)
+test_websocket_with_curl() {
+    local tenant_id="$1"
+    local job_id="$2"
+    local duration="$3"
+    
+    log_warning "WebSocket testing tools not available (websocat/node.js)"
+    echo "Falling back to HTTP polling simulation..."
+    
+    # Create a simple HTTP-based progress polling simulation
+    local poll_endpoint="$API_BASE/api/v1/tenants/$tenant_id/training/jobs/$job_id/status"
+    local end_time=$(($(date +%s) + duration))
+    local poll_count=0
+    
+    echo "Simulating real-time updates by polling: $poll_endpoint"
+    echo "Duration: ${duration}s"
+    
+    while [ $(date +%s) -lt $end_time ]; do
+        poll_count=$((poll_count + 1))
+        echo ""
+        echo "[$(date '+%H:%M:%S')] Poll #$poll_count - Checking training status..."
+        
+        STATUS_RESPONSE=$(curl -s -X GET "$poll_endpoint" \
+            -H "Authorization: Bearer $ACCESS_TOKEN" \
+            -H "X-Tenant-ID: $tenant_id")
+        
+        echo "Response:"
+        echo "$STATUS_RESPONSE" | python3 -m json.tool 2>/dev/null || echo "$STATUS_RESPONSE"
+        
+        # Check if training is complete
+        if echo "$STATUS_RESPONSE" | grep -q '"status".*"completed"\|"status".*"failed"'; then
+            log_success "Training status detected as complete/failed - stopping polling"
+            break
+        fi
+        
+        sleep 5
+    done
+    
+    log_success "HTTP polling simulation completed ($poll_count polls)"
+    echo "💡 For real WebSocket testing, install: npm install -g websocat"
+}
+
+# Wait for WebSocket messages and analyze them
+wait_for_websocket_messages() {
+    local ws_log="$1"
+    local duration="$2"
+    local start_time=$(date +%s)
+    local end_time=$((start_time + duration))
+    
+    echo "📡 Monitoring WebSocket messages..."
+    echo "Log file: $ws_log"
+    
+    # Show real-time progress
+    while [ $(date +%s) -lt $end_time ]; do
+        if [ -f "$ws_log" ]; then
+            local message_count=$(wc -l < "$ws_log" 2>/dev/null || echo "0")
+            local elapsed=$(($(date +%s) - start_time))
+            printf "\r⏱️  Elapsed: ${elapsed}s | Messages: $message_count"
+        fi
+        sleep 1
+    done
+    
+    echo ""
+    
+    # Analyze received messages
+    if [ -f "$ws_log" ] && [ -s "$ws_log" ]; then
+        local total_messages=$(wc -l < "$ws_log")
+        log_success "WebSocket test completed - received $total_messages messages"
+        
+        echo ""
+        echo "📊 Message Analysis:"
+        
+        # Show message types
+        if grep -q "progress" "$ws_log"; then
+            local progress_count=$(grep -c "progress" "$ws_log")
+            echo "  📈 Progress updates: $progress_count"
+        fi
+        
+        if grep -q "completed" "$ws_log"; then
+            echo "  ✅ Completion messages: $(grep -c "completed" "$ws_log")"
+        fi
+        
+        if grep -q "failed\|error" "$ws_log"; then
+            echo "  ❌ Error messages: $(grep -c "failed\|error" "$ws_log")"
+        fi
+        
+        echo ""
+        echo "📝 Recent messages (last 5):"
+        tail -5 "$ws_log" | sed 's/^/  /'
+        
+    else
+        log_warning "No WebSocket messages received during test period"
+        echo "  This could mean:"
+        echo "  • Training completed before WebSocket connection was established"
+        echo "  • WebSocket endpoint is not working correctly"
+        echo "  • Authentication issues with WebSocket connection"
+        echo "  • Training service is not publishing progress events"
+    fi
+    
+    # Clean up log file
+    rm -f "$ws_log"
+}
+
+# Enhanced training step with WebSocket testing
+enhanced_training_step_with_completion_check() {
+    echo -e "${STEP_ICONS[3]} ${PURPLE}STEP 4: MODEL TRAINING WITH WEBSOCKET MONITORING${NC}"
+    echo "Enhanced training step with real-time progress monitoring"
+    echo ""
+
+    log_step "4.1. Initiating model training with FULL dataset"
+
+    # Start training job
+    TRAINING_RESPONSE=$(curl -s -w "\nHTTP_CODE:%{http_code}" -X POST "$API_BASE/api/v1/tenants/$TENANT_ID/training/jobs" \
+        -H "Authorization: Bearer $ACCESS_TOKEN" \
+        -H "Content-Type: application/json" \
+        -d '{}')
+
+    # Extract HTTP code and response
+    HTTP_CODE=$(echo "$TRAINING_RESPONSE" | grep "HTTP_CODE:" | cut -d: -f2)
+    TRAINING_RESPONSE=$(echo "$TRAINING_RESPONSE" | sed '/HTTP_CODE:/d')
+
+    echo "Training HTTP Status Code: $HTTP_CODE"
+    echo "Training Response:"
+    echo "$TRAINING_RESPONSE" | python3 -m json.tool 2>/dev/null || echo "$TRAINING_RESPONSE"
+
+    if [ "$HTTP_CODE" = "200" ] || [ "$HTTP_CODE" = "201" ]; then
+        # Extract training job details
+        TRAINING_TASK_ID=$(extract_json_field "$TRAINING_RESPONSE" "task_id")
+        JOB_ID=$(extract_json_field "$TRAINING_RESPONSE" "job_id")
+        JOB_STATUS=$(extract_json_field "$TRAINING_RESPONSE" "status")
+        
+        # Use job_id if available, otherwise use task_id
+        WEBSOCKET_JOB_ID="${JOB_ID:-$TRAINING_TASK_ID}"
+        
+        if [ -n "$WEBSOCKET_JOB_ID" ]; then
+            log_success "Training job started successfully"
+            echo "  Job ID: $WEBSOCKET_JOB_ID"
+            echo "  Status: $JOB_STATUS"
+            
+            # Check if training completed instantly
+            if [ "$JOB_STATUS" = "completed" ]; then
+                log_warning "Training completed instantly - no real-time progress to monitor"
+                echo "  This can happen when:"
+                echo "    • No valid products found in sales data"
+                echo "    • Training data is insufficient"
+                echo "    • Models are already trained and cached"
+                echo ""
+                
+                # Show training results
+                TOTAL_PRODUCTS=$(extract_json_field "$TRAINING_RESPONSE" "training_results.total_products")
+                SUCCESSFUL_TRAININGS=$(extract_json_field "$TRAINING_RESPONSE" "training_results.successful_trainings")
+                SALES_RECORDS=$(extract_json_field "$TRAINING_RESPONSE" "data_summary.sales_records")
+                
+                echo "📊 Training Summary:"
+                echo "  Sales records: $SALES_RECORDS"
+                echo "  Products found: $TOTAL_PRODUCTS"
+                echo "  Successful trainings: $SUCCESSFUL_TRAININGS"
+                
+                if [ "$TOTAL_PRODUCTS" = "0" ]; then
+                    log_warning "No products found for training"
+                    echo "  Possible causes:"
+                    echo "    • CSV doesn't contain valid product names"
+                    echo "    • Product column is missing or malformed"
+                    echo "    • Insufficient sales data per product"
+                fi
+                
+                # Still test WebSocket for demonstration
+                log_step "4.2. Testing WebSocket endpoint (demonstration mode)"
+                echo "Even though training is complete, testing WebSocket connection..."
+                test_websocket_with_nodejs_builtin "$TENANT_ID" "$WEBSOCKET_JOB_ID" "10"
+                
+            else
+                # Training is in progress - monitor with WebSocket
+                log_step "4.2. Connecting to WebSocket for real-time progress monitoring"
+                test_websocket_with_nodejs_builtin "$TENANT_ID" "$WEBSOCKET_JOB_ID" "$WS_TEST_DURATION"
+            fi
+            
+        else
+            log_warning "Training started but couldn't extract job ID for WebSocket testing"
+            echo "Response: $TRAINING_RESPONSE"
+        fi
+    else
+        log_error "Training job failed to start (HTTP $HTTP_CODE)"
+        echo "Response: $TRAINING_RESPONSE"
+    fi
+
+    echo ""
+}
 # =================================================================
 # PRE-FLIGHT CHECKS
 # =================================================================
@@ -197,6 +786,41 @@ services_check() {
     done
 }
 
+check_websocket_prerequisites() {
+    echo -e "${PURPLE}🔍 Checking WebSocket testing prerequisites...${NC}"
+    
+    # Check for websocat
+    if command -v websocat >/dev/null 2>&1; then
+        log_success "websocat found - will use for WebSocket testing"
+        return 0
+    fi
+    
+    # Check for Node.js
+    if command -v node >/dev/null 2>&1; then
+        local node_version=$(node --version 2>/dev/null || echo "unknown")
+        log_success "Node.js found ($node_version) - will use for WebSocket testing"
+        
+        # Check if ws module is available (try to require it)
+        if node -e "require('ws')" 2>/dev/null; then
+            log_success "Node.js 'ws' module available"
+        else
+            log_warning "Node.js 'ws' module not found"
+            echo "  Install with: npm install -g ws"
+            echo "  Will attempt to use built-in functionality..."
+        fi
+        return 0
+    fi
+    
+    log_warning "Neither websocat nor Node.js found"
+    echo "  WebSocket testing will use HTTP polling fallback"
+    echo "  For better testing, install one of:"
+    echo "    • websocat: cargo install websocat"
+    echo "    • Node.js: https://nodejs.org/"
+    
+    return 1
+}
+
+
 services_check
 echo ""
 
@@ -543,19 +1167,9 @@ echo ""
 # STEP 4: MODEL TRAINING (ONBOARDING PAGE STEP 4)
 # =================================================================
 
-echo -e "${STEP_ICONS[3]} ${PURPLE}STEP 4: MODEL TRAINING${NC}"
+check_websocket_prerequisites
-echo "Simulating onboarding page step 4 - 'Entrenamiento del Modelo'"
-echo ""
 
-log_step "4.1. Initiating model training with FULL dataset"
+enhanced_training_step_with_completion_check
-
-TRAINING_RESPONSE=$(curl -s -X POST "$API_BASE/api/v1/tenants/$TENANT_ID/training/jobs" \
-  -H "Authorization: Bearer $ACCESS_TOKEN" \
-  -H "Content-Type: application/json" \
-  -d '{}')
-
-echo "Training Response:"
-echo "$TRAINING_RESPONSE" | python3 -m json.tool 2>/dev/null || echo "$TRAINING_RESPONSE"
 
 echo ""
 
@@ -568,7 +1182,7 @@ log_step "5.1. Testing basic dashboard functionality"
 # forecast request with proper schema
 FORECAST_REQUEST="{
     \"product_name\": \"pan\",
-    \"forecast_date\": \"2025-07-30\",
+    \"forecast_date\": \"2025-08-02\",
     \"forecast_days\": 1,
     \"location\": \"madrid_centro\",
     \"confidence_level\": 0.85