ai/bin/client-test

#!/usr/bin/env bash
set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

usage() {
    echo "Usage: $0 [COMMAND]"
    echo ""
    echo "Test client-py against the development environment"
    echo ""
    echo "Commands:"
    echo "  ohlc        Test OHLCClient API (default)"
    echo "  history     Test low-level HistoryClient"
    echo "  shell       Open Python shell with client installed"
    echo ""
    echo "Examples:"
    echo "  $0              # Run OHLC client test"
    echo "  $0 ohlc         # Run OHLC client test"
    echo "  $0 history      # Run history client test"
    echo "  $0 shell        # Interactive Python shell"
    exit 1
}

COMMAND="${1:-ohlc}"

check_kubectl() {
    if ! command -v kubectl &> /dev/null; then
        echo -e "${RED}Error: kubectl not found. Please install kubectl first.${NC}"
        exit 1
    fi
}

check_services() {
    echo -e "${BLUE}Checking if services are running...${NC}"

    # Check if required pods are running
    local services=("relay" "flink-jobmanager" "iceberg-catalog")
    local missing=()

    for service in "${services[@]}"; do
        if ! kubectl get pods -l app="$service" 2>/dev/null | grep -q "Running"; then
            missing+=("$service")
        fi
    done

    if [ ${#missing[@]} -gt 0 ]; then
        echo -e "${RED}Error: Required services not running: ${missing[*]}${NC}"
        echo -e "${YELLOW}Run 'bin/dev start' first to start the environment${NC}"
        exit 1
    fi

    echo -e "${GREEN}✓ All required services are running${NC}"
}

setup_port_forwards() {
    echo -e "${BLUE}Setting up port forwards...${NC}"

    # Kill any existing port forwards
    pkill -f "kubectl port-forward.*relay" 2>/dev/null || true
    pkill -f "kubectl port-forward.*iceberg-catalog" 2>/dev/null || true
    pkill -f "kubectl port-forward.*minio" 2>/dev/null || true

    # Port forward relay (5558=market-data pub, 5559=client requests)
    kubectl port-forward svc/relay 5558:5558 5559:5559 >/dev/null 2>&1 &
    local relay_pid=$!

    # Port forward iceberg-catalog (8181)
    kubectl port-forward svc/iceberg-catalog 8181:8181 >/dev/null 2>&1 &
    local iceberg_pid=$!

    # Port forward MinIO (9000) - needed for PyIceberg to read data files
    kubectl port-forward svc/minio 9000:9000 >/dev/null 2>&1 &
    local minio_pid=$!

    # Wait for port forwards to establish
    sleep 2

    echo -e "${GREEN}✓ Port forwards established${NC}"
    echo -e "${YELLOW}  Relay:           localhost:5558 (market-data), 5559 (requests)${NC}"
    echo -e "${YELLOW}  Iceberg Catalog: localhost:8181${NC}"
    echo -e "${YELLOW}  MinIO:           localhost:9000${NC}"

    # Store PIDs for cleanup
    export PORT_FORWARD_PIDS="$relay_pid $iceberg_pid $minio_pid"
}

cleanup_port_forwards() {
    if [ -n "$PORT_FORWARD_PIDS" ]; then
        echo -e "\n${BLUE}Cleaning up port forwards...${NC}"
        for pid in $PORT_FORWARD_PIDS; do
            kill $pid 2>/dev/null || true
        done
    fi
}

run_ohlc_test() {
    echo -e "${BLUE}Running OHLCClient test...${NC}"
    echo ""

    cd "$ROOT_DIR"

    # Install client-py in development mode
    pip install -e client-py >/dev/null 2>&1 || {
        echo -e "${YELLOW}Installing client-py dependencies...${NC}"
        pip install -e client-py
    }

    # Run the test
    python3 test/history_client/client_ohlc_api.py
}

run_history_test() {
    echo -e "${BLUE}Running HistoryClient test...${NC}"
    echo ""

    cd "$ROOT_DIR"

    # Install client-py in development mode
    pip install -e client-py >/dev/null 2>&1 || {
        echo -e "${YELLOW}Installing client-py dependencies...${NC}"
        pip install -e client-py
    }

    # Run the low-level test
    python3 test/history_client/client.py
}

open_shell() {
    echo -e "${BLUE}Opening Python shell with dexorder client...${NC}"
    echo ""

    cd "$ROOT_DIR"

    # Install client-py in development mode
    pip install -e client-py >/dev/null 2>&1 || {
        echo -e "${YELLOW}Installing client-py dependencies...${NC}"
        pip install -e client-py
    }

    echo -e "${BLUE}Example usage:${NC}"
    echo -e "  from dexorder import OHLCClient"
    echo -e "  import asyncio"
    echo -e "  client = OHLCClient('http://localhost:8181', 'tcp://localhost:5559', 'tcp://localhost:5558',"
    echo -e "                      s3_endpoint='http://localhost:9000', s3_access_key='minio', s3_secret_key='minio123')"
    echo -e "  # Use asyncio.run(client.fetch_ohlc(...)) to fetch data"
    echo ""

    python3 -i -c "
import sys
import os
sys.path.insert(0, os.path.join(os.getcwd(), 'client-py'))
from dexorder import OHLCClient, HistoryClient, IcebergClient
import asyncio
print('✓ dexorder package imported')
print('Available: OHLCClient, HistoryClient, IcebergClient, asyncio')
"
}

# Set up cleanup trap
trap cleanup_port_forwards EXIT

# Main command routing
check_kubectl

case "$COMMAND" in
    ohlc)
        check_services
        setup_port_forwards
        run_ohlc_test
        ;;
    history)
        check_services
        setup_port_forwards
        run_history_test
        ;;
    shell)
        check_services
        setup_port_forwards
        open_shell
        ;;
    -h|--help|help)
        usage
        ;;
    *)
        echo -e "${RED}Unknown command: $COMMAND${NC}"
        echo ""
        usage
        ;;
esac