added test case generation

scripts/generate_test_datasets.py

@@ -0,0 +1,411 @@
+#!/usr/bin/env python3
+"""
+PyArrow Dataset Generator for ML Inference Service
+
+Generates test datasets for academic challenges and model validation.
+Creates 100 PyArrow datasets with various image types and test scenarios.
+"""
+
+import base64
+import json
+import random
+from pathlib import Path
+from typing import Dict, List, Any, Tuple
+import io
+
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+from PIL import Image, ImageDraw, ImageFont
+
+
+class TestDatasetGenerator:
+    def __init__(self, output_dir: str = "test_datasets"):
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(exist_ok=True)
+
+        # ImageNet class labels (sample for testing)
+        self.imagenet_labels = [
+            "tench", "goldfish", "great_white_shark", "tiger_shark", "hammerhead",
+            "electric_ray", "stingray", "cock", "hen", "ostrich", "brambling",
+            "goldfinch", "house_finch", "junco", "indigo_bunting", "robin",
+            "bulbul", "jay", "magpie", "chickadee", "water_ouzel", "kite",
+            "bald_eagle", "vulture", "great_grey_owl", "European_fire_salamander",
+            "common_newt", "eft", "spotted_salamander", "axolotl", "bullfrog",
+            "tree_frog", "tailed_frog", "loggerhead", "leatherback_turtle",
+            "mud_turtle", "terrapin", "box_turtle", "banded_gecko", "common_iguana",
+            "American_chameleon", "whiptail", "agama", "frilled_lizard", "alligator_lizard",
+            "Gila_monster", "green_lizard", "African_chameleon", "Komodo_dragon",
+            "African_crocodile", "American_alligator", "triceratops", "thunder_snake"
+        ]
+
+    def create_synthetic_image(self, width: int = 224, height: int = 224,
+                             image_type: str = "random") -> Image.Image:
+        """Create synthetic images for testing."""
+        if image_type == "random":
+            # Random noise image
+            array = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
+            return Image.fromarray(array)
+
+        elif image_type == "geometric":
+            # Geometric patterns
+            img = Image.new('RGB', (width, height), color='white')
+            draw = ImageDraw.Draw(img)
+
+            # Draw random shapes
+            for _ in range(random.randint(3, 8)):
+                color = tuple(random.randint(0, 255) for _ in range(3))
+                shape_type = random.choice(['rectangle', 'ellipse'])
+                x1, y1 = random.randint(0, width//2), random.randint(0, height//2)
+                x2, y2 = x1 + random.randint(20, width//2), y1 + random.randint(20, height//2)
+
+                if shape_type == 'rectangle':
+                    draw.rectangle([x1, y1, x2, y2], fill=color)
+                else:
+                    draw.ellipse([x1, y1, x2, y2], fill=color)
+
+            return img
+
+        elif image_type == "gradient":
+            array = np.zeros((height, width, 3), dtype=np.uint8)
+            for i in range(height):
+                for j in range(width):
+                    array[i, j] = [i * 255 // height, j * 255 // width, (i + j) * 255 // (height + width)]
+            return Image.fromarray(array)
+
+        elif image_type == "text":
+            img = Image.new('RGB', (width, height), color='white')
+            draw = ImageDraw.Draw(img)
+
+            try:
+                font = ImageFont.load_default()
+            except:
+                font = None
+
+            text = f"Test Image {random.randint(1, 1000)}"
+            draw.text((width//4, height//2), text, fill='black', font=font)
+            return img
+
+        else:
+            color = tuple(random.randint(0, 255) for _ in range(3))
+            return Image.new('RGB', (width, height), color=color)
+
+    def image_to_base64(self, image: Image.Image, format: str = "JPEG") -> str:
+        """Convert PIL image to base64 string."""
+        buffer = io.BytesIO()
+        image.save(buffer, format=format)
+        image_bytes = buffer.getvalue()
+        return base64.b64encode(image_bytes).decode('utf-8')
+
+    def create_api_request(self, image_b64: str, media_type: str = "image/jpeg") -> Dict[str, Any]:
+        """Create API request structure matching your service."""
+        return {
+            "image": {
+                "mediaType": media_type,
+                "data": image_b64
+            }
+        }
+
+    def create_expected_response(self, model_name: str = "microsoft/resnet-18",
+                               media_type: str = "image/jpeg") -> Dict[str, Any]:
+        """Create expected response structure."""
+        prediction = random.choice(self.imagenet_labels)
+        return {
+            "prediction": prediction,
+            "confidence": round(random.uniform(0.3, 0.99), 4),
+            "predicted_label": random.randint(0, len(self.imagenet_labels) - 1),
+            "model": model_name,
+            "mediaType": media_type
+        }
+
+    def generate_standard_datasets(self, count: int = 25) -> List[Dict[str, Any]]:
+        """Generate standard test cases with normal images."""
+        datasets = []
+
+        for i in range(count):
+            image_types = ["random", "geometric", "gradient", "text", "solid"]
+            sizes = [(224, 224), (256, 256), (299, 299), (384, 384)]
+            formats = [("JPEG", "image/jpeg"), ("PNG", "image/png")]
+
+            records = []
+            for j in range(random.randint(5, 20)):  # 5-20 images per dataset
+                img_type = random.choice(image_types)
+                size = random.choice(sizes)
+                format_info = random.choice(formats)
+
+                image = self.create_synthetic_image(size[0], size[1], img_type)
+                image_b64 = self.image_to_base64(image, format_info[0])
+
+                api_request = self.create_api_request(image_b64, format_info[1])
+                expected_response = self.create_expected_response()
+
+                record = {
+                    "dataset_id": f"standard_{i:03d}",
+                    "image_id": f"img_{j:03d}",
+                    "image_type": img_type,
+                    "image_size": f"{size[0]}x{size[1]}",
+                    "format": format_info[0],
+                    "media_type": format_info[1],
+                    "api_request": json.dumps(api_request),
+                    "expected_response": json.dumps(expected_response),
+                    "test_category": "standard",
+                    "difficulty": "normal"
+                }
+                records.append(record)
+
+            datasets.append({
+                "name": f"standard_test_{i:03d}",
+                "category": "standard",
+                "description": f"Standard test dataset {i+1} with {len(records)} images",
+                "records": records
+            })
+
+        return datasets
+
+    def generate_edge_case_datasets(self, count: int = 25) -> List[Dict[str, Any]]:
+        """Generate datasets for edge case scenarios."""
+        datasets = []
+
+        for i in range(count):
+            records = []
+            edge_cases = [
+                {"type": "tiny", "size": (32, 32), "difficulty": "high"},
+                {"type": "huge", "size": (2048, 2048), "difficulty": "high"},
+                {"type": "extreme_aspect", "size": (1000, 50), "difficulty": "medium"},
+                {"type": "single_pixel", "size": (1, 1), "difficulty": "extreme"},
+                {"type": "corrupted_base64", "size": (224, 224), "difficulty": "extreme"}
+            ]
+
+            for j, edge_case in enumerate(edge_cases):
+                if edge_case["type"] == "corrupted_base64":
+                    image = self.create_synthetic_image(224, 224, "random")
+                    image_b64 = self.image_to_base64(image, "JPEG")
+                    corrupted_b64 = image_b64[:-20] + "CORRUPTED_DATA"
+                    api_request = self.create_api_request(corrupted_b64)
+                    expected_response = {
+                        "error": "Invalid image data",
+                        "status": "failed"
+                    }
+                else:
+                    image = self.create_synthetic_image(
+                        edge_case["size"][0], edge_case["size"][1], "random"
+                    )
+                    image_b64 = self.image_to_base64(image, "PNG")
+                    api_request = self.create_api_request(image_b64, "image/png")
+                    expected_response = self.create_expected_response()
+
+                record = {
+                    "dataset_id": f"edge_{i:03d}",
+                    "image_id": f"edge_{j:03d}",
+                    "image_type": edge_case["type"],
+                    "image_size": f"{edge_case['size'][0]}x{edge_case['size'][1]}",
+                    "format": "PNG",
+                    "media_type": "image/png",
+                    "api_request": json.dumps(api_request),
+                    "expected_response": json.dumps(expected_response),
+                    "test_category": "edge_case",
+                    "difficulty": edge_case["difficulty"]
+                }
+                records.append(record)
+
+            datasets.append({
+                "name": f"edge_case_{i:03d}",
+                "category": "edge_case",
+                "description": f"Edge case dataset {i+1} with challenging scenarios",
+                "records": records
+            })
+
+        return datasets
+
+    def generate_performance_datasets(self, count: int = 25) -> List[Dict[str, Any]]:
+        """Generate performance benchmark datasets."""
+        datasets = []
+
+        for i in range(count):
+            batch_sizes = [1, 5, 10, 25, 50, 100]
+            batch_size = random.choice(batch_sizes)
+
+            records = []
+            for j in range(batch_size):
+                image = self.create_synthetic_image(224, 224, "random")
+                image_b64 = self.image_to_base64(image, "JPEG")
+                api_request = self.create_api_request(image_b64)
+                expected_response = self.create_expected_response()
+
+                record = {
+                    "dataset_id": f"perf_{i:03d}",
+                    "image_id": f"batch_{j:03d}",
+                    "image_type": "performance_test",
+                    "image_size": "224x224",
+                    "format": "JPEG",
+                    "media_type": "image/jpeg",
+                    "api_request": json.dumps(api_request),
+                    "expected_response": json.dumps(expected_response),
+                    "test_category": "performance",
+                    "difficulty": "normal",
+                    "batch_size": batch_size,
+                    "expected_max_latency_ms": batch_size * 100
+                }
+                records.append(record)
+
+            datasets.append({
+                "name": f"performance_test_{i:03d}",
+                "category": "performance",
+                "description": f"Performance dataset {i+1} with batch size {batch_size}",
+                "records": records
+            })
+
+        return datasets
+
+    def generate_model_comparison_datasets(self, count: int = 25) -> List[Dict[str, Any]]:
+        """Generate datasets for comparing different models."""
+        datasets = []
+
+        model_types = [
+            "microsoft/resnet-18", "microsoft/resnet-50", "google/vit-base-patch16-224",
+            "facebook/convnext-tiny-224", "microsoft/swin-tiny-patch4-window7-224"
+        ]
+
+        for i in range(count):
+            # Same images tested across different model types
+            base_images = []
+            for _ in range(10):  # 10 base images per comparison dataset
+                image = self.create_synthetic_image(224, 224, "geometric")
+                base_images.append(self.image_to_base64(image, "JPEG"))
+
+            records = []
+            for j, model in enumerate(model_types):
+                for k, image_b64 in enumerate(base_images):
+                    api_request = self.create_api_request(image_b64)
+                    expected_response = self.create_expected_response(model)
+
+                    record = {
+                        "dataset_id": f"comparison_{i:03d}",
+                        "image_id": f"img_{k:03d}_model_{j}",
+                        "image_type": "comparison_base",
+                        "image_size": "224x224",
+                        "format": "JPEG",
+                        "media_type": "image/jpeg",
+                        "api_request": json.dumps(api_request),
+                        "expected_response": json.dumps(expected_response),
+                        "test_category": "model_comparison",
+                        "difficulty": "normal",
+                        "model_type": model,
+                        "comparison_group": k
+                    }
+                    records.append(record)
+
+            datasets.append({
+                "name": f"model_comparison_{i:03d}",
+                "category": "model_comparison",
+                "description": f"Model comparison dataset {i+1} testing {len(model_types)} models",
+                "records": records
+            })
+
+        return datasets
+
+    def save_dataset_to_parquet(self, dataset: Dict[str, Any]):
+        """Save a dataset to PyArrow Parquet format."""
+        records = dataset["records"]
+
+        # Convert to PyArrow table
+        table = pa.table({
+            "dataset_id": [r["dataset_id"] for r in records],
+            "image_id": [r["image_id"] for r in records],
+            "image_type": [r["image_type"] for r in records],
+            "image_size": [r["image_size"] for r in records],
+            "format": [r["format"] for r in records],
+            "media_type": [r["media_type"] for r in records],
+            "api_request": [r["api_request"] for r in records],
+            "expected_response": [r["expected_response"] for r in records],
+            "test_category": [r["test_category"] for r in records],
+            "difficulty": [r["difficulty"] for r in records],
+            # Optional fields with defaults
+            "batch_size": [r.get("batch_size", 1) for r in records],
+            "expected_max_latency_ms": [r.get("expected_max_latency_ms", 1000) for r in records],
+            "model_type": [r.get("model_type", "microsoft/resnet-18") for r in records],
+            "comparison_group": [r.get("comparison_group", 0) for r in records]
+        })
+
+        output_path = self.output_dir / f"{dataset['name']}.parquet"
+        pq.write_table(table, output_path)
+
+        # Save metadata as JSON
+        metadata = {
+            "name": dataset["name"],
+            "category": dataset["category"],
+            "description": dataset["description"],
+            "record_count": len(records),
+            "file_size_mb": round(output_path.stat().st_size / (1024 * 1024), 2),
+            "schema": [field.name for field in table.schema]
+        }
+
+        metadata_path = self.output_dir / f"{dataset['name']}_metadata.json"
+        with open(metadata_path, 'w') as f:
+            json.dump(metadata, f, indent=2)
+
+    def generate_all_datasets(self):
+        """Generate all 100 datasets."""
+        print(" Starting dataset generation...")
+
+        print("📊 Generating standard test datasets (25)...")
+        standard_datasets = self.generate_standard_datasets(25)
+        for dataset in standard_datasets:
+            self.save_dataset_to_parquet(dataset)
+
+        print("⚡ Generating edge case datasets (25)...")
+        edge_datasets = self.generate_edge_case_datasets(25)
+        for dataset in edge_datasets:
+            self.save_dataset_to_parquet(dataset)
+
+        print("🏁 Generating performance datasets (25)...")
+        performance_datasets = self.generate_performance_datasets(25)
+        for dataset in performance_datasets:
+            self.save_dataset_to_parquet(dataset)
+
+        print("🔄 Generating model comparison datasets (25)...")
+        comparison_datasets = self.generate_model_comparison_datasets(25)
+        for dataset in comparison_datasets:
+            self.save_dataset_to_parquet(dataset)
+
+        print(f"✅ Generated 100 datasets in {self.output_dir}/")
+
+        self.generate_summary()
+
+    def generate_summary(self):
+        """Generate a summary of all datasets."""
+        summary = {
+            "total_datasets": 100,
+            "categories": {
+                "standard": 25,
+                "edge_case": 25,
+                "performance": 25,
+                "model_comparison": 25
+            },
+            "dataset_info": [],
+            "usage_instructions": {
+                "loading": "Use pyarrow.parquet.read_table('dataset.parquet')",
+                "testing": "Run python scripts/test_datasets.py",
+                "api_endpoint": "POST /predict/resnet",
+                "request_format": "See api_request column in datasets"
+            }
+        }
+
+        # Add individual dataset info
+        for parquet_file in self.output_dir.glob("*.parquet"):
+            metadata_file = self.output_dir / f"{parquet_file.stem}_metadata.json"
+            if metadata_file.exists():
+                with open(metadata_file, 'r') as f:
+                    metadata = json.load(f)
+                    summary["dataset_info"].append(metadata)
+
+        summary_path = self.output_dir / "datasets_summary.json"
+        with open(summary_path, 'w') as f:
+            json.dump(summary, f, indent=2)
+
+        print(f"📋 Summary saved to {summary_path}")
+
+
+if __name__ == "__main__":
+    generator = TestDatasetGenerator()
+    generator.generate_all_datasets()

scripts/test_datasets.py

@@ -0,0 +1,382 @@
+#!/usr/bin/env python3
+"""
+Dataset Tester for ML Inference Service
+
+Tests the generated PyArrow datasets against the running ML inference service.
+Validates API requests/responses and measures performance metrics.
+"""
+
+import json
+import time
+import asyncio
+import statistics
+from pathlib import Path
+from typing import Dict, List, Any, Optional
+import argparse
+
+import pyarrow.parquet as pq
+import requests
+import pandas as pd
+
+
+class DatasetTester:
+    def __init__(self, base_url: str = "http://127.0.0.1:8000", datasets_dir: str = "test_datasets"):
+        self.base_url = base_url.rstrip('/')
+        self.datasets_dir = Path(datasets_dir)
+        self.endpoint = f"{self.base_url}/predict/resnet"
+        self.results = []
+
+    def load_dataset(self, dataset_path: Path) -> pd.DataFrame:
+        """Load a PyArrow dataset."""
+        table = pq.read_table(dataset_path)
+        return table.to_pandas()
+
+    def test_api_connection(self) -> bool:
+        """Test if the API is running and accessible."""
+        try:
+            response = requests.get(f"{self.base_url}/docs", timeout=5)
+            return response.status_code == 200
+        except requests.RequestException:
+            return False
+
+    def send_prediction_request(self, api_request_json: str) -> Dict[str, Any]:
+        """Send a single prediction request to the API."""
+        try:
+            request_data = json.loads(api_request_json)
+            start_time = time.time()
+
+            response = requests.post(
+                self.endpoint,
+                json=request_data,
+                headers={"Content-Type": "application/json"},
+                timeout=30
+            )
+
+            end_time = time.time()
+            latency_ms = (end_time - start_time) * 1000
+
+            return {
+                "success": response.status_code == 200,
+                "status_code": response.status_code,
+                "response": response.json() if response.status_code == 200 else response.text,
+                "latency_ms": round(latency_ms, 2),
+                "error": None
+            }
+
+        except requests.RequestException as e:
+            return {
+                "success": False,
+                "status_code": None,
+                "response": None,
+                "latency_ms": None,
+                "error": str(e)
+            }
+        except json.JSONDecodeError as e:
+            return {
+                "success": False,
+                "status_code": None,
+                "response": None,
+                "latency_ms": None,
+                "error": f"JSON decode error: {str(e)}"
+            }
+
+    def validate_response(self, actual_response: Dict[str, Any],
+                         expected_response_json: str) -> Dict[str, Any]:
+        """Validate API response against expected response."""
+        try:
+            expected = json.loads(expected_response_json)
+
+            validation = {
+                "structure_valid": True,
+                "field_errors": []
+            }
+
+            # Check required fields exist
+            required_fields = ["prediction", "confidence", "predicted_label", "model", "mediaType"]
+            for field in required_fields:
+                if field not in actual_response:
+                    validation["structure_valid"] = False
+                    validation["field_errors"].append(f"Missing field: {field}")
+
+            # Validate field types
+            if "confidence" in actual_response:
+                if not isinstance(actual_response["confidence"], (int, float)):
+                    validation["field_errors"].append("confidence must be numeric")
+                elif not (0 <= actual_response["confidence"] <= 1):
+                    validation["field_errors"].append("confidence must be between 0 and 1")
+
+            if "predicted_label" in actual_response:
+                if not isinstance(actual_response["predicted_label"], int):
+                    validation["field_errors"].append("predicted_label must be integer")
+
+            return validation
+
+        except json.JSONDecodeError:
+            return {
+                "structure_valid": False,
+                "field_errors": ["Invalid expected response JSON"]
+            }
+
+    def test_dataset(self, dataset_path: Path, max_samples: Optional[int] = None) -> Dict[str, Any]:
+        """Test a single dataset."""
+        print(f"📊 Testing dataset: {dataset_path.name}")
+
+        try:
+            df = self.load_dataset(dataset_path)
+            if max_samples:
+                df = df.head(max_samples)
+
+            results = {
+                "dataset_name": dataset_path.stem,
+                "total_samples": len(df),
+                "tested_samples": 0,
+                "successful_requests": 0,
+                "failed_requests": 0,
+                "validation_errors": 0,
+                "latencies_ms": [],
+                "errors": [],
+                "category": df['test_category'].iloc[0] if not df.empty else "unknown"
+            }
+
+            for idx, row in df.iterrows():
+                print(f"  Testing sample {idx + 1}/{len(df)}", end="\r")
+
+                # Send API request
+                api_result = self.send_prediction_request(row['api_request'])
+                results["tested_samples"] += 1
+
+                if api_result["success"]:
+                    results["successful_requests"] += 1
+                    results["latencies_ms"].append(api_result["latency_ms"])
+
+                    # Validate response structure
+                    validation = self.validate_response(
+                        api_result["response"],
+                        row['expected_response']
+                    )
+
+                    if not validation["structure_valid"]:
+                        results["validation_errors"] += 1
+                        results["errors"].append({
+                            "sample_id": row['image_id'],
+                            "type": "validation_error",
+                            "details": validation["field_errors"]
+                        })
+
+                else:
+                    results["failed_requests"] += 1
+                    results["errors"].append({
+                        "sample_id": row['image_id'],
+                        "type": "request_failed",
+                        "status_code": api_result["status_code"],
+                        "error": api_result["error"]
+                    })
+
+            # Calculate statistics
+            if results["latencies_ms"]:
+                results["avg_latency_ms"] = round(statistics.mean(results["latencies_ms"]), 2)
+                results["min_latency_ms"] = round(min(results["latencies_ms"]), 2)
+                results["max_latency_ms"] = round(max(results["latencies_ms"]), 2)
+                results["median_latency_ms"] = round(statistics.median(results["latencies_ms"]), 2)
+            else:
+                results.update({
+                    "avg_latency_ms": None,
+                    "min_latency_ms": None,
+                    "max_latency_ms": None,
+                    "median_latency_ms": None
+                })
+
+            results["success_rate"] = round(
+                results["successful_requests"] / results["tested_samples"] * 100, 2
+            ) if results["tested_samples"] > 0 else 0
+
+            print(f"\n  ✅ Completed: {results['success_rate']}% success rate")
+            return results
+
+        except Exception as e:
+            print(f"\n  ❌ Failed to test dataset: {str(e)}")
+            return {
+                "dataset_name": dataset_path.stem,
+                "error": str(e),
+                "success_rate": 0
+            }
+
+    def test_all_datasets(self, max_samples_per_dataset: Optional[int] = None,
+                         category_filter: Optional[str] = None) -> Dict[str, Any]:
+        """Test all datasets or filtered by category."""
+        if not self.test_api_connection():
+            print("❌ API is not accessible. Please start the service first:")
+            print("   uvicorn main:app --reload")
+            return {"error": "API not accessible"}
+
+        print(f" Starting dataset testing against {self.endpoint}")
+
+        parquet_files = list(self.datasets_dir.glob("*.parquet"))
+        if not parquet_files:
+            print(f"❌ No datasets found in {self.datasets_dir}")
+            return {"error": "No datasets found"}
+
+        if category_filter:
+            parquet_files = [f for f in parquet_files if category_filter in f.name]
+
+        print(f" Found {len(parquet_files)} datasets to test")
+
+        all_results = []
+        start_time = time.time()
+
+        for dataset_file in parquet_files:
+            result = self.test_dataset(dataset_file, max_samples_per_dataset)
+            all_results.append(result)
+
+        end_time = time.time()
+        total_time = end_time - start_time
+
+        summary = self.generate_summary(all_results, total_time)
+
+        self.save_results(summary, all_results)
+
+        return summary
+
+    def generate_summary(self, results: List[Dict[str, Any]], total_time: float) -> Dict[str, Any]:
+        """Generate summary of all test results."""
+        successful_datasets = [r for r in results if r.get("success_rate", 0) > 0]
+        failed_datasets = [r for r in results if r.get("error") or r.get("success_rate", 0) == 0]
+
+        total_samples = sum(r.get("tested_samples", 0) for r in results)
+        total_successful = sum(r.get("successful_requests", 0) for r in results)
+        total_failed = sum(r.get("failed_requests", 0) for r in results)
+
+        all_latencies = []
+        for r in results:
+            all_latencies.extend(r.get("latencies_ms", []))
+
+        summary = {
+            "test_summary": {
+                "total_datasets": len(results),
+                "successful_datasets": len(successful_datasets),
+                "failed_datasets": len(failed_datasets),
+                "total_samples_tested": total_samples,
+                "total_successful_requests": total_successful,
+                "total_failed_requests": total_failed,
+                "overall_success_rate": round(
+                    total_successful / total_samples * 100, 2
+                ) if total_samples > 0 else 0,
+                "total_test_time_seconds": round(total_time, 2)
+            },
+            "performance_metrics": {
+                "avg_latency_ms": round(statistics.mean(all_latencies), 2) if all_latencies else None,
+                "median_latency_ms": round(statistics.median(all_latencies), 2) if all_latencies else None,
+                "min_latency_ms": round(min(all_latencies), 2) if all_latencies else None,
+                "max_latency_ms": round(max(all_latencies), 2) if all_latencies else None,
+                "requests_per_second": round(
+                    total_successful / total_time, 2
+                ) if total_time > 0 else 0
+            },
+            "category_breakdown": {},
+            "failed_datasets": [r["dataset_name"] for r in failed_datasets]
+        }
+
+        categories = {}
+        for result in results:
+            category = result.get("category", "unknown")
+            if category not in categories:
+                categories[category] = {
+                    "count": 0,
+                    "success_rates": [],
+                    "avg_success_rate": 0
+                }
+            categories[category]["count"] += 1
+            categories[category]["success_rates"].append(result.get("success_rate", 0))
+
+        for category, data in categories.items():
+            data["avg_success_rate"] = round(
+                statistics.mean(data["success_rates"]), 2
+            ) if data["success_rates"] else 0
+
+        summary["category_breakdown"] = categories
+
+        return summary
+
+    def save_results(self, summary: Dict[str, Any], detailed_results: List[Dict[str, Any]]):
+        """Save test results to files."""
+        results_dir = Path("test_results")
+        results_dir.mkdir(exist_ok=True)
+
+        timestamp = int(time.time())
+
+        # Save summary
+        summary_path = results_dir / f"test_summary_{timestamp}.json"
+        with open(summary_path, 'w') as f:
+            json.dump(summary, f, indent=2)
+
+        # Save detailed results
+        detailed_path = results_dir / f"test_detailed_{timestamp}.json"
+        with open(detailed_path, 'w') as f:
+            json.dump(detailed_results, f, indent=2)
+
+        print(f" Results saved:")
+        print(f"   Summary: {summary_path}")
+        print(f"   Details: {detailed_path}")
+
+    def print_summary(self, summary: Dict[str, Any]):
+        """Print test summary to console."""
+        print("\n" + "="*60)
+        print("🏁 DATASET TESTING SUMMARY")
+        print("="*60)
+
+        ts = summary["test_summary"]
+        print(f"Datasets tested: {ts['total_datasets']}")
+        print(f"Successful datasets: {ts['successful_datasets']}")
+        print(f"Failed datasets: {ts['failed_datasets']}")
+        print(f"Total samples: {ts['total_samples_tested']}")
+        print(f"Overall success rate: {ts['overall_success_rate']}%")
+        print(f"Test duration: {ts['total_test_time_seconds']}s")
+
+        pm = summary["performance_metrics"]
+        if pm["avg_latency_ms"]:
+            print(f"\nPerformance:")
+            print(f"  Avg latency: {pm['avg_latency_ms']}ms")
+            print(f"  Median latency: {pm['median_latency_ms']}ms")
+            print(f"  Min latency: {pm['min_latency_ms']}ms")
+            print(f"  Max latency: {pm['max_latency_ms']}ms")
+            print(f"  Requests/sec: {pm['requests_per_second']}")
+
+        print(f"\nCategory breakdown:")
+        for category, data in summary["category_breakdown"].items():
+            print(f"  {category}: {data['count']} datasets, {data['avg_success_rate']}% avg success")
+
+        if summary["failed_datasets"]:
+            print(f"\nFailed datasets: {', '.join(summary['failed_datasets'])}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Test PyArrow datasets against ML inference service")
+    parser.add_argument("--base-url", default="http://127.0.0.1:8000", help="Base URL of the API")
+    parser.add_argument("--datasets-dir", default="scripts/test_datasets", help="Directory containing datasets")
+    parser.add_argument("--max-samples", type=int, help="Max samples per dataset to test")
+    parser.add_argument("--category", help="Filter datasets by category (standard, edge_case, performance, model_comparison)")
+    parser.add_argument("--quick", action="store_true", help="Quick test with max 5 samples per dataset")
+
+    args = parser.parse_args()
+
+    tester = DatasetTester(args.base_url, args.datasets_dir)
+
+    max_samples = args.max_samples
+    if args.quick:
+        max_samples = 5
+
+    results = tester.test_all_datasets(max_samples, args.category)
+
+    if "error" not in results:
+        tester.print_summary(results)
+
+        if results["test_summary"]["overall_success_rate"] > 90:
+            print("\n🎉 Excellent! API is working great with the datasets!")
+        elif results["test_summary"]["overall_success_rate"] > 70:
+            print("\n👍 Good! API works well, minor issues detected.")
+        else:
+            print("\n⚠️  Warning: Several issues detected. Check the detailed results.")
+
+
+if __name__ == "__main__":
+    main()