diff --git a/README.md b/README.md index d4c9131..633a530 100644 --- a/README.md +++ b/README.md @@ -12,8 +12,8 @@ This library is designed to be integrated into robot data collection systems, ET - 🤖 **Teleoperation Session Management**: Track robot data collection sessions with `TeleopSession` - 🔄 **Data Conversion Session Management**: Track data conversion pipelines with `ConversionSession` -- ☁️ **S3 Upload Management**: Track S3 object storage uploads with `S3UploadSession` -- 🔌 **USB Data Copy Management**: Track USB data copy operations with `USBCopySession` +- ☁️ **S3 Upload Management**: Track S3 object storage uploads with `S3UploadSession` (includes device and job tracking) +- 🔌 **USB Data Copy Management**: Track USB data copy operations with `USBCopySession` (includes device and job tracking) - 📊 **OpenLineage Integration**: Full OpenLineage specification support (START/COMPLETE/RUNNING/FAIL/ABORT) - 🏷️ **Custom Run Facets**: Robot metadata (robotId, location, repository info) via `CommonRunFacet` - 🔍 **Marquez Client**: Query jobs, datasets, and lineage graphs via REST API @@ -124,11 +124,16 @@ Track USB copy operations from the command line. Quick example: ```bash +# Generate unique job ID +JOB_ID=$(python3 -c "import uuid; print(uuid.uuid4())") + # Start USB copy session RUN_ID=$(airoa-lineage-usb-copy start \ --namespace production \ --robot-id hsr001 \ --location weblab \ + --hostname $(hostname) \ + --job-id "$JOB_ID" \ --repository-hash $(git rev-parse HEAD) \ --repository-uri https://github.com/user/repo.git \ --repository-tag v1.0.0 \ @@ -140,7 +145,7 @@ RUN_ID=$(airoa-lineage-usb-copy start \ # ... # Complete session -airoa-lineage-usb-copy complete --run-id "$RUN_ID" +airoa-lineage-usb-copy complete --run-id "$RUN_ID" --job-id "$JOB_ID" --hostname $(hostname) ``` See [docs/cli-usage.md](docs/cli-usage.md) for full CLI reference. diff --git a/docs/architecture.md b/docs/architecture.md index 98e5a66..82a4415 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -103,6 +103,10 @@ ConversionSession.complete(output_datasets) # Emits COMPLETE event - Track S3 data upload sessions (START → COMPLETE) - Automatic run_id generation and management - Nominal time support for historical data processing +- Device tracking via DeviceRunFacet (hostname - **required**) +- Job tracking via JobRunFacet (unique job ID - **required**) +- Input datasets tracking (source directory, source files) +- Output datasets tracking (destination path, upload duration) - State validation (prevents duplicate start/complete calls) **Lifecycle**: @@ -125,6 +129,10 @@ S3UploadSession.complete() # Emits COMPLETE event - Track USB data copy sessions (START → COMPLETE) - Automatic run_id generation and management - Nominal time support for historical data processing +- Device tracking via DeviceRunFacet (hostname - **required**) +- Job tracking via JobRunFacet (unique job ID - **required**) +- Input datasets tracking (source directory, source files) +- Output datasets tracking (destination directory, copy duration) - State validation (prevents duplicate start/complete calls) **Lifecycle**: @@ -158,13 +166,18 @@ USBCopySession.complete() # Emits COMPLETE event **Purpose**: Provide robot-specific metadata extensions to OpenLineage Run Facets. **Available Facets**: -- **CommonRunFacet** ([common.py](../src/airoa_lineage/facets/common.py)): Robot metadata (robotId, location, repository info) -- **DeviceRunFacet** ([device.py](../src/airoa_lineage/facets/device.py)): Device metadata (hostname) -- **AWSJobRunFacet** ([aws_job.py](../src/airoa_lineage/facets/aws_job.py)): AWS Lambda job metadata (name, id) +- **CommonRunFacet** ([run/common.py](../src/airoa_lineage/facets/run/common.py)): Robot metadata (robotId, location, repository info) + - Used by: TeleopSession, ConversionSession, S3UploadSession, USBCopySession +- **DeviceRunFacet** ([run/device.py](../src/airoa_lineage/facets/run/device.py)): Device metadata (hostname) + - Used by: TeleopSession, S3UploadSession (**required**), USBCopySession (**required**) +- **JobRunFacet** ([run/job.py](../src/airoa_lineage/facets/run/job.py)): Job identification metadata (unique job ID) + - Used by: S3UploadSession (**required**), USBCopySession (**required**) +- **AWSJobRunFacet** ([run/aws_job.py](../src/airoa_lineage/facets/run/aws_job.py)): AWS Lambda job metadata (name, id) + - Used by: ConversionSession -**Facet Namespacing**: Use `facet_prefix` parameter to customize facet names (e.g., `facet_prefix="airoa"` → facets named `"airoa_common"`, `"airoa_device"`, and `"airoa_awsJob"`) +**Facet Namespacing**: Use `facet_prefix` parameter to customize facet names (e.g., `facet_prefix="airoa"` → facets named `"airoa_common"`, `"airoa_device"`, `"airoa_job"`, and `"airoa_awsJob"`) -**API Documentation**: See docstrings in [facets/common.py](../src/airoa_lineage/facets/common.py) and [facets/device.py](../src/airoa_lineage/facets/device.py) +**API Documentation**: See docstrings in [facets/run/common.py](../src/airoa_lineage/facets/run/common.py), [facets/run/device.py](../src/airoa_lineage/facets/run/device.py), and [facets/run/job.py](../src/airoa_lineage/facets/run/job.py) ## How Components Work Together @@ -281,9 +294,10 @@ For detailed event schemas, see the [OpenLineage specification](https://openline - [S3UploadSession](../src/airoa_lineage/s3_upload/session.py) - S3 object storage upload lifecycle - [USBCopySession](../src/airoa_lineage/usb_copy/session.py) - USB data copy lifecycle - [Timestamp Utilities](../src/airoa_lineage/utils/timestamps.py) - UTC timestamp generation -- [CommonRunFacet](../src/airoa_lineage/facets/common.py) - Robot metadata facet -- [DeviceRunFacet](../src/airoa_lineage/facets/device.py) - Device metadata facet -- [AWSJobRunFacet](../src/airoa_lineage/facets/aws_job.py) - AWS Lambda job metadata facet +- [CommonRunFacet](../src/airoa_lineage/facets/run/common.py) - Robot metadata facet +- [DeviceRunFacet](../src/airoa_lineage/facets/run/device.py) - Device metadata facet +- [JobRunFacet](../src/airoa_lineage/facets/run/job.py) - Job identification metadata facet +- [AWSJobRunFacet](../src/airoa_lineage/facets/run/aws_job.py) - AWS Lambda job metadata facet ### Examples (Usage Patterns) diff --git a/docs/cli/usb-copy.md b/docs/cli/usb-copy.md index d244c72..5cede14 100644 --- a/docs/cli/usb-copy.md +++ b/docs/cli/usb-copy.md @@ -21,12 +21,16 @@ airoa-lineage-usb-copy start [OPTIONS] | `--namespace` | string | Yes* | OpenLineage namespace | | `--robot-id` | string | Yes* | Robot identifier | | `--location` | string | Yes* | Location identifier | +| `--hostname` | string | Yes* | Hostname of the device performing the USB copy | +| `--job-id` | string | Yes* | Unique job identifier (e.g., UUID) | | `--repository-hash` | string | Yes* | Git commit hash | | `--repository-uri` | string | Yes* | Repository URI | | `--repository-tag` | string | Yes* | Git tag | | `--repository-branch` | string | Yes* | Git branch | | `--nominal-start-time` | string | No | Nominal start time (ISO 8601 format) | | `--nominal-end-time` | string | No | Nominal end time (ISO 8601 format) | +| `--input-dataset` | string | No | JSON string containing input dataset metadata (mutually exclusive with --input-dataset-file) | +| `--input-dataset-file` | string | No | Path to JSON file containing input dataset metadata (mutually exclusive with --input-dataset) | | `--job-name` | string | No | Job name (default: `usb-data-copy`) | | `--marquez-url` | string | No | Marquez server URL | | `--facet-prefix` | string | No | Facet prefix | @@ -49,6 +53,8 @@ RUN_ID=$(airoa-lineage-usb-copy start \ --namespace airoa_examples \ --robot-id hsr001 \ --location weblab \ + --hostname $(hostname) \ + --job-id $(python3 -c "import uuid; print(uuid.uuid4())") \ --repository-hash $(git rev-parse HEAD) \ --repository-uri https://github.com/AIRoA/airoa-lineage.git \ --repository-tag v1.0.0 \ @@ -56,6 +62,7 @@ RUN_ID=$(airoa-lineage-usb-copy start \ # With nominal times (for batch processing) RUN_ID=$(airoa-lineage-usb-copy start \ + --job-id $(python3 -c "import uuid; print(uuid.uuid4())") \ --nominal-start-time "2025-11-26T00:00:00+00:00" \ --nominal-end-time "2025-11-26T05:00:00+00:00") @@ -89,10 +96,14 @@ airoa-lineage-usb-copy complete --run-id RUN_ID [OPTIONS] | `--facet-prefix` | string | No | Facet prefix | | `--robot-id` | string | Yes* | Robot identifier | | `--location` | string | Yes* | Location identifier | +| `--hostname` | string | Yes* | Hostname of the device performing the USB copy | +| `--job-id` | string | Yes* | Unique job identifier (same as used in start) | | `--repository-hash` | string | Yes* | Git commit hash | | `--repository-uri` | string | Yes* | Repository URI | | `--repository-tag` | string | Yes* | Git tag | | `--repository-branch` | string | Yes* | Git branch | +| `--output-dataset` | string | No | JSON string containing output dataset metadata (mutually exclusive with --output-dataset-file) | +| `--output-dataset-file` | string | No | Path to JSON file containing output dataset metadata (mutually exclusive with --output-dataset) | \* Required unless provided via config file or environment variables @@ -107,16 +118,19 @@ airoa-lineage-usb-copy complete --run-id RUN_ID [OPTIONS] **Examples:** ```bash -# Basic usage +# Basic usage (with environment variables or config file) airoa-lineage-usb-copy complete --run-id "$RUN_ID" # Specify all arguments explicitly (using same values as start) +JOB_ID="your-job-id-from-start" # Must be the same as used in start airoa-lineage-usb-copy complete \ --run-id "$RUN_ID" \ --namespace airoa_examples \ --job-name usb-data-copy \ --robot-id hsr001 \ --location weblab \ + --hostname $(hostname) \ + --job-id "$JOB_ID" \ --repository-hash abc123 \ --repository-uri https://github.com/AIRoA/airoa-lineage.git \ --repository-tag v1.0.0 \ @@ -152,6 +166,8 @@ airoa-lineage-usb-copy cancel --run-id RUN_ID [OPTIONS] | `--facet-prefix` | string | No | Facet prefix | | `--robot-id` | string | Yes* | Robot identifier | | `--location` | string | Yes* | Location identifier | +| `--hostname` | string | Yes* | Hostname of the device performing the USB copy | +| `--job-id` | string | Yes* | Unique job identifier (same as used in start) | | `--repository-hash` | string | Yes* | Git commit hash | | `--repository-uri` | string | Yes* | Repository URI | | `--repository-tag` | string | Yes* | Git tag | @@ -170,16 +186,19 @@ airoa-lineage-usb-copy cancel --run-id RUN_ID [OPTIONS] **Examples:** ```bash -# Basic usage +# Basic usage (with environment variables or config file) airoa-lineage-usb-copy cancel --run-id "$RUN_ID" # Specify all arguments explicitly (using same values as start) +JOB_ID="your-job-id-from-start" # Must be the same as used in start airoa-lineage-usb-copy cancel \ --run-id "$RUN_ID" \ --namespace airoa_examples \ --job-name usb-data-copy \ --robot-id hsr001 \ --location weblab \ + --hostname $(hostname) \ + --job-id "$JOB_ID" \ --repository-hash abc123 \ --repository-uri https://github.com/AIRoA/airoa-lineage.git \ --repository-tag v1.0.0 \ @@ -194,6 +213,203 @@ airoa-lineage-usb-copy cancel --run-id "$RUN_ID" --quiet --- +## Dataset Tracking + +The USBCopy CLI supports tracking detailed input and output dataset metadata in two ways: +1. **JSON files** using `--input-dataset-file` and `--output-dataset-file` +2. **JSON strings** using `--input-dataset` and `--output-dataset` (useful for dynamic generation without creating temporary files) + +This allows you to record: + +- **Input datasets**: USB device information, source directories, and file lists with MD5 hashes +- **Output datasets**: Destination directories and copy operation statistics (duration) + +### Input Dataset Format + +Both `--input-dataset-file` (file path) and `--input-dataset` (JSON string) accept the following structure: + +```json +{ + "dataset_name": "usb_input_2025_12_08", + "usb_device": { + "id": "/dev/sdb1", + "label": "ROBOT_DATA_001", + "fs_type": "ext4" + }, + "src_dir": { + "path": "/media/usb0/robot_data/2025-12-08" + }, + "src_files": [ + { + "name": "rosbag_001.bag", + "hash": "5d41402abc4b2a76b9719d911017c592" + }, + { + "name": "rosbag_002.bag", + "hash": "7d793037a0760186574b0282f2f435e7" + } + ] +} +``` + +**Fields:** + +- `dataset_name` (required): Unique name for the input dataset +- `usb_device` (optional): USB device metadata + - `id`: System filesystem ID (e.g., `/dev/sdb1`) + - `label`: USB disk label (e.g., `ROBOT_DATA_001`) + - `fs_type`: Filesystem type (e.g., `ext4`, `vfat`, `exfat`) +- `src_dir` (optional): Source directory information + - `path`: Copy source directory path +- `src_files` (optional): List of source files + - `name`: File name + - `hash`: MD5 hash value for the file + +### Output Dataset Format + +Both `--output-dataset-file` (file path) and `--output-dataset` (JSON string) accept the following structure: + +```json +{ + "dataset_name": "usb_output_2025_12_08", + "dest_dir": { + "path": "/data/robot_data/2025-12-08" + }, + "operation_stats": { + "duration_seconds": 125.47 + } +} +``` + +**Fields:** + +- `dataset_name` (required): Unique name for the output dataset +- `dest_dir` (optional): Destination directory information + - `path`: Copy destination directory path +- `operation_stats` (optional): Operation statistics + - `duration_seconds`: Duration of the operation in seconds + +### Measuring Copy Duration + +The recommended approach for measuring copy duration is to record timestamps in your shell script before and after the copy operation: + +```bash +# Record start time +START_TIME=$(date +%s) + +# Perform copy operation +# ... your copy commands here ... + +# Record end time +END_TIME=$(date +%s) + +# Calculate duration +DURATION=$((END_TIME - START_TIME)) + +# Create output manifest with duration +cat > /tmp/output_manifest.json < int: + return run_cli(_cli_config) +""" + +from __future__ import annotations + +import argparse +import json +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Protocol + +from openlineage.client.run import Dataset + +from airoa_lineage.cli.common_args import ( + add_common_args, + add_common_facet_args, + add_device_facet_args, + add_job_facet_args, + add_repository_args, + merge_config_with_args, +) +from airoa_lineage.cli.config import ( + build_common_facet, + build_device_facet, + build_job_facet, + load_config, +) +from airoa_lineage.cli.dataset_utils import parse_dataset_json +from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet, JobRunFacet + +# Type aliases for callback functions +InputDatasetBuilder = Callable[[Dict[str, Any], str], List[Dataset]] +OutputDatasetBuilder = Callable[[Dict[str, Any], str], List[Dataset]] + + +class CLISessionProtocol(Protocol): + """Protocol for CLI session classes. + + This defines the interface that session classes must implement + to be used with the CLI framework. Both USBCopySession and + S3UploadSession implement this protocol. + """ + + def __init__( + self, + namespace: str, + common_facet: CommonRunFacet, + device_facet: DeviceRunFacet, + job_facet: JobRunFacet, + job_name: str = ..., + marquez_url: Optional[str] = ..., + run_id: Optional[str] = ..., + facet_prefix: str = ..., + ) -> None: ... + + def start( + self, + input_datasets: List[Dataset], + nominal_start_time: Optional[str] = ..., + nominal_end_time: Optional[str] = ..., + ) -> str: ... + + def resume(self, run_id: str) -> None: ... + + def complete(self, output_datasets: List[Dataset]) -> None: ... + + def cancel(self) -> None: ... + + +@dataclass(frozen=True) +class CLIConfig: + """Configuration for a CLI application. + + This dataclass encapsulates all CLI-specific parameters and callbacks, + enabling shared functions to handle different CLI implementations. + + Attributes: + prog: Program name (e.g., "airoa-lineage-usb-copy") + description: CLI description for help text + version: Version string (e.g., "airoa-lineage-usb-copy 0.1.0") + default_job_name: Default job name if not specified + session_class: Session class to instantiate (e.g., USBCopySession) + build_input_dataset: Callback to build input datasets from JSON data + build_output_dataset: Callback to build output datasets from JSON data + start_help: Help text for start subcommand + complete_help: Help text for complete subcommand + cancel_help: Help text for cancel subcommand + + Example: + >>> from airoa_lineage.usb_copy import USBCopySession + >>> config = CLIConfig( + ... prog="airoa-lineage-usb-copy", + ... description="Track USB data copy operations with OpenLineage", + ... version="airoa-lineage-usb-copy 0.1.0", + ... default_job_name="usb-data-copy", + ... session_class=USBCopySession, + ... build_input_dataset=build_usb_input_dataset, + ... build_output_dataset=build_usb_output_dataset, + ... ) + """ + + prog: str + description: str + version: str + default_job_name: str + session_class: type[CLISessionProtocol] + build_input_dataset: InputDatasetBuilder + build_output_dataset: OutputDatasetBuilder + start_help: str = "Start session and emit START event" + complete_help: str = "Complete session and emit COMPLETE event" + cancel_help: str = "Cancel session and emit ABORT event" + + +def create_parser(cli_config: CLIConfig) -> argparse.ArgumentParser: + """ + Create argument parser for CLI. + + Args: + cli_config: CLI configuration + + Returns: + Configured ArgumentParser instance + """ + parser = argparse.ArgumentParser( + prog=cli_config.prog, + description=cli_config.description, + ) + + parser.add_argument( + "--version", + action="version", + version=cli_config.version, + ) + + parser.add_argument( + "--config", + type=str, + help="Path to config file (default: ~/.config/airoa-lineage/config.json)", + ) + + parser.add_argument( + "-v", + "--verbose", + action="store_true", + help="Enable verbose output", + ) + + parser.add_argument( + "-q", + "--quiet", + action="store_true", + help="Suppress output (except errors)", + ) + + parser.add_argument( + "--json", + action="store_true", + help="Output in JSON format", + ) + + # Create subparsers for commands + subparsers = parser.add_subparsers(dest="command", required=True) + + # START command + start_parser = subparsers.add_parser("start", help=cli_config.start_help) + _add_start_args(start_parser) + + # COMPLETE command + complete_parser = subparsers.add_parser("complete", help=cli_config.complete_help) + _add_complete_args(complete_parser) + + # CANCEL command + cancel_parser = subparsers.add_parser("cancel", help=cli_config.cancel_help) + _add_cancel_args(cancel_parser) + + return parser + + +def _add_start_args(parser: argparse.ArgumentParser) -> None: + """Add arguments for start command.""" + add_common_args(parser) + add_common_facet_args(parser) + add_repository_args(parser) + add_device_facet_args(parser) + add_job_facet_args(parser) + parser.add_argument( + "--nominal-start-time", + help="Nominal start time (ISO 8601 format)", + ) + parser.add_argument( + "--nominal-end-time", + help="Nominal end time (ISO 8601 format)", + ) + # Input dataset options (mutually exclusive) + input_group = parser.add_mutually_exclusive_group() + input_group.add_argument( + "--input-dataset", + type=str, + help="JSON string containing input dataset metadata", + ) + input_group.add_argument( + "--input-dataset-file", + type=str, + help="Path to JSON file containing input dataset metadata", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show what would be done without executing", + ) + + +def _add_complete_args(parser: argparse.ArgumentParser) -> None: + """Add arguments for complete command.""" + parser.add_argument( + "--run-id", + required=True, + help="Session run ID (from start command)", + ) + add_common_args(parser) + add_common_facet_args(parser) + add_repository_args(parser) + add_device_facet_args(parser) + add_job_facet_args(parser) + # Output dataset options (mutually exclusive) + output_group = parser.add_mutually_exclusive_group() + output_group.add_argument( + "--output-dataset", + type=str, + help="JSON string containing output dataset metadata", + ) + output_group.add_argument( + "--output-dataset-file", + type=str, + help="Path to JSON file containing output dataset metadata", + ) + + +def _add_cancel_args(parser: argparse.ArgumentParser) -> None: + """Add arguments for cancel command.""" + parser.add_argument( + "--run-id", + required=True, + help="Session run ID (from start command)", + ) + add_common_args(parser) + add_common_facet_args(parser) + add_repository_args(parser) + add_device_facet_args(parser) + add_job_facet_args(parser) + + +def load_dataset_from_file( + file_path: str, + namespace: str, + builder: Callable[[Dict[str, Any], str], List[Dataset]], + dataset_type: str, +) -> List[Dataset]: + """ + Load dataset from JSON manifest file. + + Args: + file_path: Path to JSON manifest file + namespace: OpenLineage namespace + builder: Function to build datasets from parsed data + dataset_type: Type description for error messages ("input" or "output") + + Returns: + List of Dataset objects + + Raises: + ValueError: If file cannot be read or has invalid format + """ + try: + with open(file_path, "r") as f: + data = json.load(f) + except FileNotFoundError: + raise ValueError( + f"{dataset_type.capitalize()} dataset file not found: {file_path}" + ) + except json.JSONDecodeError as e: + raise ValueError(f"Invalid JSON in {dataset_type} dataset file: {e}") + + return builder(data, namespace) + + +def cmd_start( + args: argparse.Namespace, + config: Dict[str, Any], + cli_config: CLIConfig, +) -> int: + """ + Execute start command. + + Args: + args: Parsed CLI arguments + config: Merged configuration + cli_config: CLI configuration + + Returns: + Exit code (0=success, 1=error, 3=config error) + """ + try: + # Build facets + common_facet = build_common_facet(config) + device_facet = build_device_facet(config) + job_facet = build_job_facet(config) + + # Validate required fields + if "namespace" not in config: + print( + "Error: namespace is required. " + "Provide via --namespace, config file, or AIROA_NAMESPACE env var.", + file=sys.stderr, + ) + return 1 + + # Dry run mode + if args.dry_run: + _print_dry_run_info( + config, common_facet, device_facet, job_facet, cli_config, args + ) + return 0 + + # Load input datasets + input_datasets = _load_input_datasets(args, config, cli_config) + + # Create session + session = cli_config.session_class( + namespace=config["namespace"], + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + job_name=config.get("job_name", cli_config.default_job_name), + marquez_url=config.get("marquez_url"), + facet_prefix=config.get("facet_prefix", ""), + ) + + # Start session + run_id = session.start( + input_datasets=input_datasets, + nominal_start_time=getattr(args, "nominal_start_time", None), + nominal_end_time=getattr(args, "nominal_end_time", None), + ) + + # Output run_id + _print_start_output(args, run_id) + return 0 + + except ValueError as e: + print(f"Configuration error: {e}", file=sys.stderr) + return 3 + except RuntimeError as e: + print(f"Session error: {e}", file=sys.stderr) + return 1 + except Exception as e: + print(f"Unexpected error: {e}", file=sys.stderr) + if args.verbose: + import traceback + + traceback.print_exc() + return 1 + + +def cmd_complete( + args: argparse.Namespace, + config: Dict[str, Any], + cli_config: CLIConfig, +) -> int: + """ + Execute complete command. + + Args: + args: Parsed CLI arguments + config: Merged configuration + cli_config: CLI configuration + + Returns: + Exit code (0=success, 1=error, 3=config error) + """ + try: + # Build facets + common_facet = build_common_facet(config) + device_facet = build_device_facet(config) + job_facet = build_job_facet(config) + + # Validate required fields + if "namespace" not in config: + print( + "Error: namespace is required. " + "Provide via --namespace, config file, or AIROA_NAMESPACE env var.", + file=sys.stderr, + ) + return 1 + + # Load output datasets + output_datasets = _load_output_datasets(args, config, cli_config) + + # Create and resume session + session = cli_config.session_class( + namespace=config["namespace"], + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + job_name=config.get("job_name", cli_config.default_job_name), + marquez_url=config.get("marquez_url"), + facet_prefix=config.get("facet_prefix", ""), + ) + session.resume(args.run_id) + session.complete(output_datasets=output_datasets) + + # Output success + _print_complete_output(args) + return 0 + + except ValueError as e: + print(f"Configuration error: {e}", file=sys.stderr) + return 3 + except RuntimeError as e: + print(f"Session error: {e}", file=sys.stderr) + return 1 + except Exception as e: + print(f"Unexpected error: {e}", file=sys.stderr) + if args.verbose: + import traceback + + traceback.print_exc() + return 1 + + +def cmd_cancel( + args: argparse.Namespace, + config: Dict[str, Any], + cli_config: CLIConfig, +) -> int: + """ + Execute cancel command. + + Args: + args: Parsed CLI arguments + config: Merged configuration + cli_config: CLI configuration + + Returns: + Exit code (0=success, 1=error, 3=config error) + """ + try: + # Build facets + common_facet = build_common_facet(config) + device_facet = build_device_facet(config) + job_facet = build_job_facet(config) + + # Validate required fields + if "namespace" not in config: + print( + "Error: namespace is required. " + "Provide via --namespace, config file, or AIROA_NAMESPACE env var.", + file=sys.stderr, + ) + return 1 + + # Create and resume session + session = cli_config.session_class( + namespace=config["namespace"], + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + job_name=config.get("job_name", cli_config.default_job_name), + marquez_url=config.get("marquez_url"), + facet_prefix=config.get("facet_prefix", ""), + ) + session.resume(args.run_id) + session.cancel() + + # Output success + _print_cancel_output(args) + return 0 + + except ValueError as e: + print(f"Configuration error: {e}", file=sys.stderr) + return 3 + except RuntimeError as e: + print(f"Session error: {e}", file=sys.stderr) + return 1 + except Exception as e: + print(f"Unexpected error: {e}", file=sys.stderr) + if args.verbose: + import traceback + + traceback.print_exc() + return 1 + + +def run_cli(cli_config: CLIConfig) -> int: + """ + Main entry point for CLI applications. + + This function orchestrates the CLI workflow: + 1. Parse command line arguments + 2. Load and merge configuration + 3. Execute the appropriate command + + Args: + cli_config: CLI configuration + + Returns: + Exit code (0=success, 1=general error, 2=connection error, 3=config error) + + Example: + >>> if __name__ == "__main__": + ... sys.exit(run_cli(usb_copy_config)) + """ + parser = create_parser(cli_config) + args = parser.parse_args() + + # Load configuration + try: + config_path = Path(args.config) if args.config else None + config = load_config(config_path) + except FileNotFoundError as e: + print(f"Configuration file not found: {e}", file=sys.stderr) + return 3 + except json.JSONDecodeError as e: + print(f"Invalid JSON in configuration file: {e}", file=sys.stderr) + return 3 + except Exception as e: + print(f"Error loading configuration: {e}", file=sys.stderr) + return 3 + + # Merge with CLI arguments + config = merge_config_with_args(config, args) + + # Execute command + if args.command == "start": + return cmd_start(args, config, cli_config) + elif args.command == "complete": + return cmd_complete(args, config, cli_config) + elif args.command == "cancel": + return cmd_cancel(args, config, cli_config) + else: + print(f"Unknown command: {args.command}", file=sys.stderr) + return 1 + + +# Helper functions for output and dataset loading + + +def _print_dry_run_info( + config: Dict[str, Any], + common_facet: Any, + device_facet: Any, + job_facet: Any, + cli_config: CLIConfig, + args: argparse.Namespace, +) -> None: + """Print dry-run information.""" + print("[DRY RUN] Would start session with:") + print(f" Namespace: {config['namespace']}") + print(f" Job name: {config.get('job_name', cli_config.default_job_name)}") + print(f" Robot ID: {common_facet.robotId}") + print(f" Location: {common_facet.location}") + print(f" Hostname: {device_facet.hostname}") + print(f" Job ID: {job_facet.id}") + print(f" Repository: {common_facet.repositoryUri}") + print(f" Commit: {common_facet.repositoryHash}") + print(f" Branch: {common_facet.repositoryBranch}") + print(f" Tag: {common_facet.repositoryTag}") + if hasattr(args, "nominal_start_time") and args.nominal_start_time: + print(f" Nominal start: {args.nominal_start_time}") + if hasattr(args, "nominal_end_time") and args.nominal_end_time: + print(f" Nominal end: {args.nominal_end_time}") + print("[DRY RUN] Would output run_id (not generated in dry-run mode)") + + +def _load_input_datasets( + args: argparse.Namespace, + config: Dict[str, Any], + cli_config: CLIConfig, +) -> List[Dataset]: + """Load input datasets from args.""" + if hasattr(args, "input_dataset") and args.input_dataset: + data = parse_dataset_json(args.input_dataset, "input") + return cli_config.build_input_dataset(data, config["namespace"]) + elif hasattr(args, "input_dataset_file") and args.input_dataset_file: + return load_dataset_from_file( + args.input_dataset_file, + config["namespace"], + cli_config.build_input_dataset, + "input", + ) + return [] + + +def _load_output_datasets( + args: argparse.Namespace, + config: Dict[str, Any], + cli_config: CLIConfig, +) -> List[Dataset]: + """Load output datasets from args.""" + if hasattr(args, "output_dataset") and args.output_dataset: + data = parse_dataset_json(args.output_dataset, "output") + return cli_config.build_output_dataset(data, config["namespace"]) + elif hasattr(args, "output_dataset_file") and args.output_dataset_file: + return load_dataset_from_file( + args.output_dataset_file, + config["namespace"], + cli_config.build_output_dataset, + "output", + ) + return [] + + +def _print_start_output(args: argparse.Namespace, run_id: str) -> None: + """Print start command output.""" + if args.json: + output = {"run_id": run_id, "status": "started"} + print(json.dumps(output)) + elif args.verbose: + print("Session started successfully", file=sys.stderr) + print(run_id) + else: + print(run_id) + + +def _print_complete_output(args: argparse.Namespace) -> None: + """Print complete command output.""" + if not args.quiet: + if args.json: + output = {"run_id": args.run_id, "status": "completed"} + print(json.dumps(output)) + else: + print(f"Session {args.run_id} completed successfully") + + +def _print_cancel_output(args: argparse.Namespace) -> None: + """Print cancel command output.""" + if not args.quiet: + if args.json: + output = {"run_id": args.run_id, "status": "cancelled"} + print(json.dumps(output)) + else: + print(f"Session {args.run_id} cancelled successfully") diff --git a/src/airoa_lineage/cli/common_args.py b/src/airoa_lineage/cli/common_args.py new file mode 100644 index 0000000..0c7ec73 --- /dev/null +++ b/src/airoa_lineage/cli/common_args.py @@ -0,0 +1,175 @@ +"""Common CLI argument definitions for airoa-lineage commands. + +This module provides reusable argument parser functions for CLI commands +in the airoa-lineage package. These functions add standardized arguments +for OpenLineage tracking, facets, and configuration. +""" + +from __future__ import annotations + +import argparse +from typing import Any, Dict, List, Tuple + +# Mapping of CLI argument names to top-level config keys +_TOP_LEVEL_ARGS: List[Tuple[str, str]] = [ + ("namespace", "namespace"), + ("marquez_url", "marquez_url"), + ("job_name", "job_name"), + ("facet_prefix", "facet_prefix"), +] + +# Mapping of CLI argument names to nested config keys (arg_name, section, key) +_NESTED_ARGS: List[Tuple[str, str, str]] = [ + ("robot_id", "common_facet", "robotId"), + ("location", "common_facet", "location"), + ("repository_hash", "common_facet", "repositoryHash"), + ("repository_uri", "common_facet", "repositoryUri"), + ("repository_tag", "common_facet", "repositoryTag"), + ("repository_branch", "common_facet", "repositoryBranch"), + ("hostname", "device_facet", "hostname"), + ("job_id", "job_facet", "id"), +] + + +def add_common_args(parser: argparse.ArgumentParser) -> None: + """Add common arguments shared across all commands. + + These arguments can be provided via CLI, environment variables, + or configuration file (CLI takes highest precedence). + + Args: + parser: ArgumentParser instance to add arguments to + """ + parser.add_argument( + "--namespace", + help="OpenLineage namespace (can be set via MARQUEZ_NAMESPACE env var)", + ) + parser.add_argument( + "--job-name", + help="Job name (can be set via JOB_NAME env var)", + ) + parser.add_argument( + "--marquez-url", + help="Marquez server URL (can be set via MARQUEZ_URL env var)", + ) + parser.add_argument( + "--facet-prefix", + default="", + help="Facet prefix (default: empty string)", + ) + + +def add_common_facet_args(parser: argparse.ArgumentParser) -> None: + """Add CommonRunFacet arguments. + + These arguments represent robot and location metadata required + for OpenLineage event tracking. + + Args: + parser: ArgumentParser instance to add arguments to + """ + parser.add_argument( + "--robot-id", + help="Robot identifier (e.g., 'hsr001')", + ) + parser.add_argument( + "--location", + help="Location identifier (e.g., 'weblab', 'factory-floor-2')", + ) + + +def add_repository_args(parser: argparse.ArgumentParser) -> None: + """Add repository-related arguments for CommonRunFacet. + + These arguments represent Git repository metadata for tracking + code versions in OpenLineage events. + + Args: + parser: ArgumentParser instance to add arguments to + """ + parser.add_argument( + "--repository-hash", + help="Git commit hash", + ) + parser.add_argument( + "--repository-uri", + help="Repository URI (e.g., 'https://github.com/org/repo.git')", + ) + parser.add_argument( + "--repository-tag", + help="Git tag (e.g., 'v1.0.0')", + ) + parser.add_argument( + "--repository-branch", + help="Git branch (e.g., 'main', 'develop')", + ) + + +def add_device_facet_args(parser: argparse.ArgumentParser) -> None: + """Add DeviceRunFacet arguments. + + These arguments represent device metadata required for OpenLineage event tracking. + + Args: + parser: ArgumentParser instance to add arguments to + """ + parser.add_argument( + "--hostname", + help="Device hostname (can be set via HOSTNAME env var)", + ) + + +def add_job_facet_args(parser: argparse.ArgumentParser) -> None: + """Add JobRunFacet arguments. + + These arguments represent job identification metadata required + for OpenLineage event tracking. + + Args: + parser: ArgumentParser instance to add arguments to + """ + parser.add_argument( + "--job-id", + help="Unique job identifier (UUID format recommended)", + ) + + +def merge_config_with_args( + config: Dict[str, Any], + args: argparse.Namespace, +) -> Dict[str, Any]: + """Merge configuration with CLI arguments. + + CLI arguments take precedence over config file and environment variables. + Uses mapping tables to reduce cyclomatic complexity. + + Args: + config: Configuration from load_config() + args: Parsed CLI arguments + + Returns: + Merged configuration dictionary + + Example: + >>> import argparse + >>> config = {"namespace": "default"} + >>> args = argparse.Namespace(namespace="cli_namespace", robot_id="hsr001") + >>> result = merge_config_with_args(config, args) + >>> result["namespace"] + 'cli_namespace' + """ + # Process top-level arguments + for arg_name, config_key in _TOP_LEVEL_ARGS: + value = getattr(args, arg_name, None) + if value: + config[config_key] = value + + # Process nested arguments + for arg_name, section, key in _NESTED_ARGS: + value = getattr(args, arg_name, None) + if value: + if section not in config: + config[section] = {} + config[section][key] = value + + return config diff --git a/src/airoa_lineage/cli/config.py b/src/airoa_lineage/cli/config.py index 655132f..004a876 100644 --- a/src/airoa_lineage/cli/config.py +++ b/src/airoa_lineage/cli/config.py @@ -7,7 +7,7 @@ from pathlib import Path from typing import Any, Dict, Optional -from airoa_lineage.facets import CommonRunFacet +from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet, JobRunFacet def get_default_config_path() -> Path: @@ -111,6 +111,24 @@ def load_config(config_path: Optional[Path] = None) -> Dict[str, Any]: if common_facet: config["common_facet"] = common_facet + # DeviceRunFacet fields from environment variables + device_facet = config.get("device_facet", {}) + + if hostname := os.getenv("AIROA_HOSTNAME"): + device_facet["hostname"] = hostname + + if device_facet: + config["device_facet"] = device_facet + + # JobRunFacet fields from environment variables + job_facet = config.get("job_facet", {}) + + if job_id := os.getenv("AIROA_JOB_ID"): + job_facet["id"] = job_id + + if job_facet: + config["job_facet"] = job_facet + return config @@ -168,3 +186,71 @@ def build_common_facet(config: Dict[str, Any]) -> CommonRunFacet: repositoryTag=facet_data["repositoryTag"], repositoryBranch=facet_data["repositoryBranch"], ) + + +def build_device_facet(config: Dict[str, Any]) -> DeviceRunFacet: + """ + Build DeviceRunFacet from configuration. + + Args: + config: Configuration dictionary (from load_config()) + + Returns: + DeviceRunFacet instance + + Raises: + ValueError: If required fields are missing + + Examples: + >>> config = { + ... "device_facet": { + ... "hostname": "copy-pc-001" + ... } + ... } + >>> facet = build_device_facet(config) + >>> print(facet.hostname) + copy-pc-001 + """ + facet_data = config.get("device_facet", {}) + + if "hostname" not in facet_data: + raise ValueError( + "Missing required DeviceRunFacet field: hostname. " + "Please provide it via --hostname, AIROA_HOSTNAME environment variable, or config file." + ) + + return DeviceRunFacet(hostname=facet_data["hostname"]) + + +def build_job_facet(config: Dict[str, Any]) -> JobRunFacet: + """ + Build JobRunFacet from configuration. + + Args: + config: Configuration dictionary (from load_config()) + + Returns: + JobRunFacet instance + + Raises: + ValueError: If required fields are missing + + Examples: + >>> config = { + ... "job_facet": { + ... "id": "a1b2c3d4-e5f6-7890-abcd-ef1234567890" + ... } + ... } + >>> facet = build_job_facet(config) + >>> print(facet.id) + a1b2c3d4-e5f6-7890-abcd-ef1234567890 + """ + facet_data = config.get("job_facet", {}) + + if "id" not in facet_data: + raise ValueError( + "Missing required JobRunFacet field: id. " + "Please provide it via --job-id, AIROA_JOB_ID environment variable, or config file." + ) + + return JobRunFacet(id=facet_data["id"]) diff --git a/src/airoa_lineage/cli/dataset_utils.py b/src/airoa_lineage/cli/dataset_utils.py new file mode 100644 index 0000000..336aa58 --- /dev/null +++ b/src/airoa_lineage/cli/dataset_utils.py @@ -0,0 +1,149 @@ +"""Common dataset utilities for airoa-lineage CLI commands. + +This module provides reusable functions for parsing and building OpenLineage +datasets from JSON input. These utilities are used by various CLI commands +(S3 upload, USB copy, etc.) to construct input and output datasets. +""" + +from __future__ import annotations + +import json +from typing import Any, Dict + +from openlineage.client.run import Dataset + +from airoa_lineage.facets.dataset import ( + DestDirDatasetFacet, + OperationStatsDatasetFacet, + SrcDirDatasetFacet, + SrcFileInfo, + SrcFilesDatasetFacet, +) + + +def parse_dataset_json(json_string: str, dataset_type: str) -> Dict[str, Any]: + """Parse dataset JSON string. + + Args: + json_string: JSON string to parse + dataset_type: Type description for error messages (e.g., "input", "output") + + Returns: + Parsed dataset data dictionary + + Raises: + ValueError: If JSON is invalid + + Example: + >>> data = parse_dataset_json('{"dataset_name": "test"}', "input") + >>> data["dataset_name"] + 'test' + """ + try: + data = json.loads(json_string) + return data + except json.JSONDecodeError as e: + raise ValueError(f"Invalid JSON in {dataset_type} dataset string: {e}") + + +def build_input_dataset( + data: Dict[str, Any], namespace: str, facet_prefix: str = "" +) -> Dataset: + """Build input dataset from parsed JSON data. + + Constructs an OpenLineage Dataset with appropriate facets based on the + provided data. Supports source directory and source files facets. + + Args: + data: Dataset data dictionary containing: + - dataset_name (str): Name of the dataset (required) + - src_dir (dict, optional): Source directory with 'path' key + - src_files (list, optional): List of files with 'name' and 'hash' keys + namespace: OpenLineage namespace + facet_prefix: Prefix for custom facets (default: empty string) + + Returns: + OpenLineage Dataset with appropriate facets + + Example: + >>> data = { + ... "dataset_name": "input_2025_12_10", + ... "src_dir": {"path": "/data/source"}, + ... "src_files": [ + ... {"name": "file1.txt", "hash": "abc123"}, + ... {"name": "file2.txt", "hash": "def456"} + ... ] + ... } + >>> dataset = build_input_dataset(data, "my_namespace") + >>> dataset.name + 'input_2025_12_10' + """ + # Extract dataset name + dataset_name = data.get("dataset_name", "input") + + # Build facets + facets: Dict[str, Any] = {} + + # Add SrcDirDatasetFacet if src_dir is present + if "src_dir" in data: + src_dir_data = data["src_dir"] + facets[f"{facet_prefix}srcDir"] = SrcDirDatasetFacet(path=src_dir_data["path"]) + + # Add SrcFilesDatasetFacet if src_files is present + if "src_files" in data: + src_files = [ + SrcFileInfo(name=f["name"], hash=f["hash"]) for f in data["src_files"] + ] + facets[f"{facet_prefix}srcFiles"] = SrcFilesDatasetFacet(files=src_files) + + # Create dataset + dataset = Dataset(namespace=namespace, name=dataset_name, facets=facets) + + return dataset + + +def build_output_dataset_common_facets( + data: Dict[str, Any], facet_prefix: str = "" +) -> Dict[str, Any]: + """Build common output dataset facets. + + Extracts facets that are common across different output types (S3, USB, etc.). + Operation-specific facets should be added by the caller. + + Args: + data: Dataset data dictionary containing: + - dest_dir (dict, optional): Destination directory with 'path' key + - operation_stats (dict, optional): Operation statistics with 'duration_seconds' + facet_prefix: Prefix for custom facets (default: empty string) + + Returns: + Dictionary of facets + + Example: + >>> data = { + ... "dest_dir": {"path": "/data/destination"}, + ... "operation_stats": {"duration_seconds": 125.47} + ... } + >>> facets = build_output_dataset_common_facets(data) + >>> "destDir" in facets + True + >>> "operationStats" in facets + True + """ + facets: Dict[str, Any] = {} + + # Add DestDirDatasetFacet if dest_dir is present + if "dest_dir" in data: + dest_dir_data = data["dest_dir"] + facets[f"{facet_prefix}destDir"] = DestDirDatasetFacet( + path=dest_dir_data["path"] + ) + + # Add OperationStatsDatasetFacet + if "operation_stats" in data: + stats_data = data["operation_stats"] + facets[f"{facet_prefix}operationStats"] = OperationStatsDatasetFacet( + durationSeconds=stats_data["duration_seconds"] + ) + + return facets diff --git a/src/airoa_lineage/cli/s3_upload.py b/src/airoa_lineage/cli/s3_upload.py index f4bbaa9..7f839d5 100644 --- a/src/airoa_lineage/cli/s3_upload.py +++ b/src/airoa_lineage/cli/s3_upload.py @@ -3,6 +3,9 @@ This module provides a command-line interface for tracking S3 data upload operations using OpenLineage. +This module uses common CLI utilities from airoa_lineage.cli for argument +parsing and dataset construction. + Usage: airoa-lineage-s3-upload start [options] airoa-lineage-s3-upload complete --run-id RUN_ID [options] @@ -11,451 +14,106 @@ from __future__ import annotations -import argparse -import json import sys -from pathlib import Path -from typing import Any, Dict - -from airoa_lineage.cli.config import build_common_facet, load_config -from airoa_lineage.s3_upload import S3UploadSession - - -def _add_common_args(parser: argparse.ArgumentParser) -> None: - """ - Add common arguments shared across all commands. +from typing import Any, Dict, List - These arguments can be provided via CLI, environment variables, - or configuration file (CLI takes highest precedence). +from openlineage.client.run import Dataset - Args: - parser: ArgumentParser to add arguments to - """ - parser.add_argument( - "--namespace", - help="OpenLineage namespace", - ) - parser.add_argument( - "--job-name", - help="Job name", - ) - parser.add_argument( - "--marquez-url", - help="Marquez server URL", - ) - parser.add_argument( - "--facet-prefix", - help="Facet prefix (default: empty string)", - ) +from airoa_lineage.cli.base import CLIConfig, run_cli +from airoa_lineage.cli.base import create_parser as _create_parser +from airoa_lineage.cli.common_args import merge_config_with_args # noqa: F401 (re-export) +from airoa_lineage.cli.dataset_utils import ( + build_input_dataset, + build_output_dataset_common_facets, +) +from airoa_lineage.s3_upload import S3UploadSession -def _add_common_facet_args(parser: argparse.ArgumentParser) -> None: +def _build_input_dataset(data: Dict[str, Any], namespace: str) -> List[Dataset]: """ - Add CommonRunFacet arguments. + Build input dataset for S3 upload. - These arguments represent robot and repository metadata required - for OpenLineage event tracking. + Expected data format: + { + "dataset_name": "s3_input_2025_12_08", + "src_dir": { + "path": "/data/robot_data/2025-12-08" + }, + "src_files": [ + {"name": "rosbag_001.bag", "hash": "abc123"}, + {"name": "rosbag_002.bag", "hash": "def456"} + ] + } Args: - parser: ArgumentParser to add arguments to - """ - parser.add_argument( - "--robot-id", - help="Robot identifier", - ) - parser.add_argument( - "--location", - help="Location identifier", - ) - parser.add_argument( - "--repository-hash", - help="Git commit hash", - ) - parser.add_argument( - "--repository-uri", - help="Repository URI", - ) - parser.add_argument( - "--repository-tag", - help="Git tag", - ) - parser.add_argument( - "--repository-branch", - help="Git branch", - ) - - -def create_parser() -> argparse.ArgumentParser: - """ - Create argument parser for airoa-lineage-s3-upload CLI. + data: Parsed JSON data + namespace: OpenLineage namespace Returns: - Configured ArgumentParser instance + List containing single input Dataset """ - parser = argparse.ArgumentParser( - prog="airoa-lineage-s3-upload", - description="Track S3 upload operations with OpenLineage", - ) - - parser.add_argument( - "--version", - action="version", - version="airoa-lineage-s3-upload 0.1.0", - ) - - parser.add_argument( - "--config", - type=str, - help="Path to config file (default: ~/.config/airoa-lineage/config.json)", - ) - - parser.add_argument( - "-v", - "--verbose", - action="store_true", - help="Enable verbose output", - ) - - parser.add_argument( - "-q", - "--quiet", - action="store_true", - help="Suppress output (except errors)", - ) - - parser.add_argument( - "--json", - action="store_true", - help="Output in JSON format", - ) + # S3 upload uses common input dataset structure + return [build_input_dataset(data, namespace)] - # Create subparsers for commands - subparsers = parser.add_subparsers(dest="command", required=True) - # START command - start_parser = subparsers.add_parser( - "start", - help="Start S3 upload session and emit START event", - ) - _add_common_args(start_parser) - _add_common_facet_args(start_parser) - start_parser.add_argument( - "--nominal-start-time", - help="Nominal start time (ISO 8601 format)", - ) - start_parser.add_argument( - "--nominal-end-time", - help="Nominal end time (ISO 8601 format)", - ) - start_parser.add_argument( - "--dry-run", - action="store_true", - help="Show what would be done without executing", - ) - - # COMPLETE command - complete_parser = subparsers.add_parser( - "complete", - help="Complete S3 upload session and emit COMPLETE event", - ) - complete_parser.add_argument( - "--run-id", - required=True, - help="Session run ID (from start command)", - ) - _add_common_args(complete_parser) - _add_common_facet_args(complete_parser) - - # CANCEL command - cancel_parser = subparsers.add_parser( - "cancel", - help="Cancel S3 upload session and emit ABORT event", - ) - cancel_parser.add_argument( - "--run-id", - required=True, - help="Session run ID (from start command)", - ) - _add_common_args(cancel_parser) - _add_common_facet_args(cancel_parser) - - return parser - - -def merge_config_with_args( - config: Dict[str, Any], - args: argparse.Namespace, -) -> Dict[str, Any]: +def _build_output_dataset(data: Dict[str, Any], namespace: str) -> List[Dataset]: """ - Merge configuration with CLI arguments. + Build output dataset for S3 upload. - CLI arguments take precedence over config file and environment variables. + Expected data format: + { + "dataset_name": "s3_output_2025_12_08", + "dest_dir": { + "path": "s3://my-bucket/robot_data/2025-12-08" + }, + "operation_stats": { + "duration_seconds": 325.89 + } + } Args: - config: Configuration from load_config() - args: Parsed CLI arguments + data: Parsed JSON data + namespace: OpenLineage namespace Returns: - Merged configuration dictionary + List containing single output Dataset """ - # Top-level configuration - if args.namespace: - config["namespace"] = args.namespace + # Extract dataset name + dataset_name = data.get("dataset_name", "s3_output") - if hasattr(args, "marquez_url") and args.marquez_url: - config["marquez_url"] = args.marquez_url + # Get common facets + facets = build_output_dataset_common_facets(data) - if hasattr(args, "job_name") and args.job_name: - config["job_name"] = args.job_name + # Add S3-specific facets here if needed in the future - if hasattr(args, "facet_prefix") and args.facet_prefix: - config["facet_prefix"] = args.facet_prefix + # Create dataset + dataset = Dataset(namespace=namespace, name=dataset_name, facets=facets) - # CommonRunFacet fields - if hasattr(args, "robot_id") and args.robot_id: - if "common_facet" not in config: - config["common_facet"] = {} - config["common_facet"]["robotId"] = args.robot_id + return [dataset] - if hasattr(args, "location") and args.location: - if "common_facet" not in config: - config["common_facet"] = {} - config["common_facet"]["location"] = args.location - if hasattr(args, "repository_hash") and args.repository_hash: - if "common_facet" not in config: - config["common_facet"] = {} - config["common_facet"]["repositoryHash"] = args.repository_hash +# CLI Configuration +_cli_config = CLIConfig( + prog="airoa-lineage-s3-upload", + description="Track S3 upload operations with OpenLineage", + version="airoa-lineage-s3-upload 0.1.0", + default_job_name="s3-data-upload", + session_class=S3UploadSession, + build_input_dataset=_build_input_dataset, + build_output_dataset=_build_output_dataset, + start_help="Start S3 upload session and emit START event", + complete_help="Complete S3 upload session and emit COMPLETE event", + cancel_help="Cancel S3 upload session and emit ABORT event", +) - if hasattr(args, "repository_uri") and args.repository_uri: - if "common_facet" not in config: - config["common_facet"] = {} - config["common_facet"]["repositoryUri"] = args.repository_uri - if hasattr(args, "repository_tag") and args.repository_tag: - if "common_facet" not in config: - config["common_facet"] = {} - config["common_facet"]["repositoryTag"] = args.repository_tag - - if hasattr(args, "repository_branch") and args.repository_branch: - if "common_facet" not in config: - config["common_facet"] = {} - config["common_facet"]["repositoryBranch"] = args.repository_branch - - return config - - -def cmd_start(args: argparse.Namespace, config: Dict[str, Any]) -> int: - """ - Execute start command. - - Args: - args: Parsed CLI arguments - config: Merged configuration +def create_parser(): + """Create argument parser for airoa-lineage-s3-upload CLI. Returns: - Exit code (0=success, 1=error) - """ - try: - # Build CommonRunFacet - common_facet = build_common_facet(config) - - # Validate required fields - if "namespace" not in config: - print( - "Error: namespace is required. " - "Provide via --namespace, config file, or AIROA_NAMESPACE env var.", - file=sys.stderr, - ) - return 1 - - # Dry run mode - if args.dry_run: - print("[DRY RUN] Would start session with:") - print(f" Namespace: {config['namespace']}") - print(f" Job name: {config.get('job_name', 's3-data-upload')}") - print(f" Robot ID: {common_facet.robotId}") - print(f" Location: {common_facet.location}") - print(f" Repository: {common_facet.repositoryUri}") - print(f" Commit: {common_facet.repositoryHash}") - print(f" Branch: {common_facet.repositoryBranch}") - print(f" Tag: {common_facet.repositoryTag}") - if hasattr(args, "nominal_start_time") and args.nominal_start_time: - print(f" Nominal start: {args.nominal_start_time}") - if hasattr(args, "nominal_end_time") and args.nominal_end_time: - print(f" Nominal end: {args.nominal_end_time}") - print("[DRY RUN] Would output run_id (not generated in dry-run mode)") - return 0 - - # Create session - session = S3UploadSession( - namespace=config["namespace"], - common_facet=common_facet, - job_name=config.get("job_name", "s3-data-upload"), - marquez_url=config.get("marquez_url"), - facet_prefix=config.get("facet_prefix", ""), - ) - - # Start session - run_id = session.start( - nominal_start_time=getattr(args, "nominal_start_time", None), - nominal_end_time=getattr(args, "nominal_end_time", None), - ) - - # Output run_id - if args.json: - output = {"run_id": run_id, "status": "started"} - print(json.dumps(output)) - elif args.verbose: - print("Session started successfully", file=sys.stderr) - print(run_id) - else: - # Default: just print run_id - print(run_id) - - return 0 - - except ValueError as e: - print(f"Configuration error: {e}", file=sys.stderr) - return 3 - except RuntimeError as e: - print(f"Session error: {e}", file=sys.stderr) - return 1 - except Exception as e: - print(f"Unexpected error: {e}", file=sys.stderr) - if args.verbose: - import traceback - - traceback.print_exc() - return 1 - - -def cmd_complete(args: argparse.Namespace, config: Dict[str, Any]) -> int: - """ - Execute complete command. - - Args: - args: Parsed CLI arguments - config: Merged configuration - - Returns: - Exit code (0=success, 1=error) - """ - try: - # Build CommonRunFacet - common_facet = build_common_facet(config) - - # Validate required fields - if "namespace" not in config: - print( - "Error: namespace is required. " - "Provide via --namespace, config file, or AIROA_NAMESPACE env var.", - file=sys.stderr, - ) - return 1 - - # Create session - session = S3UploadSession( - namespace=config["namespace"], - common_facet=common_facet, - job_name=config.get("job_name", "s3-data-upload"), - marquez_url=config.get("marquez_url"), - facet_prefix=config.get("facet_prefix", ""), - ) - - # Resume existing session - session.resume(args.run_id) - - # Complete session - session.complete() - - # Output success message - if not args.quiet: - if args.json: - output = {"run_id": args.run_id, "status": "completed"} - print(json.dumps(output)) - else: - print(f"Session {args.run_id} completed successfully") - - return 0 - - except ValueError as e: - print(f"Configuration error: {e}", file=sys.stderr) - return 3 - except RuntimeError as e: - print(f"Session error: {e}", file=sys.stderr) - return 1 - except Exception as e: - print(f"Unexpected error: {e}", file=sys.stderr) - if args.verbose: - import traceback - - traceback.print_exc() - return 1 - - -def cmd_cancel(args: argparse.Namespace, config: Dict[str, Any]) -> int: - """ - Execute cancel command. - - Args: - args: Parsed CLI arguments - config: Merged configuration - - Returns: - Exit code (0=success, 1=error) + Configured ArgumentParser instance """ - try: - # Build CommonRunFacet - common_facet = build_common_facet(config) - - # Validate required fields - if "namespace" not in config: - print( - "Error: namespace is required. " - "Provide via --namespace, config file, or AIROA_NAMESPACE env var.", - file=sys.stderr, - ) - return 1 - - # Create session - session = S3UploadSession( - namespace=config["namespace"], - common_facet=common_facet, - job_name=config.get("job_name", "s3-data-upload"), - marquez_url=config.get("marquez_url"), - facet_prefix=config.get("facet_prefix", ""), - ) - - # Resume existing session - session.resume(args.run_id) - - # Cancel session - session.cancel() - - # Output success message - if not args.quiet: - if args.json: - output = {"run_id": args.run_id, "status": "cancelled"} - print(json.dumps(output)) - else: - print(f"Session {args.run_id} cancelled successfully") - - return 0 - - except ValueError as e: - print(f"Configuration error: {e}", file=sys.stderr) - return 3 - except RuntimeError as e: - print(f"Session error: {e}", file=sys.stderr) - return 1 - except Exception as e: - print(f"Unexpected error: {e}", file=sys.stderr) - if args.verbose: - import traceback - - traceback.print_exc() - return 1 + return _create_parser(_cli_config) def main() -> int: @@ -465,36 +123,7 @@ def main() -> int: Returns: Exit code (0=success, 1=general error, 2=connection error, 3=config error) """ - parser = create_parser() - args = parser.parse_args() - - # Load configuration - try: - config_path = Path(args.config) if args.config else None - config = load_config(config_path) - except FileNotFoundError as e: - print(f"Configuration file not found: {e}", file=sys.stderr) - return 3 - except json.JSONDecodeError as e: - print(f"Invalid JSON in configuration file: {e}", file=sys.stderr) - return 3 - except Exception as e: - print(f"Error loading configuration: {e}", file=sys.stderr) - return 3 - - # Merge with CLI arguments - config = merge_config_with_args(config, args) - - # Execute command - if args.command == "start": - return cmd_start(args, config) - elif args.command == "complete": - return cmd_complete(args, config) - elif args.command == "cancel": - return cmd_cancel(args, config) - else: - print(f"Unknown command: {args.command}", file=sys.stderr) - return 1 + return run_cli(_cli_config) if __name__ == "__main__": diff --git a/src/airoa_lineage/cli/usb_copy.py b/src/airoa_lineage/cli/usb_copy.py index 33e6c31..acfda6e 100644 --- a/src/airoa_lineage/cli/usb_copy.py +++ b/src/airoa_lineage/cli/usb_copy.py @@ -3,6 +3,9 @@ This module provides a command-line interface for tracking USB data copy operations using OpenLineage. +This module uses common CLI utilities from airoa_lineage.cli for argument +parsing and dataset construction. + Usage: airoa-lineage-usb-copy start [options] airoa-lineage-usb-copy complete --run-id RUN_ID [options] @@ -11,451 +14,135 @@ from __future__ import annotations -import argparse -import json import sys -from pathlib import Path -from typing import Any, Dict - -from airoa_lineage.cli.config import build_common_facet, load_config +from typing import Any, Dict, List + +from openlineage.client.run import Dataset + +from airoa_lineage.cli.base import CLIConfig, run_cli +from airoa_lineage.cli.base import create_parser as _create_parser +from airoa_lineage.cli.common_args import merge_config_with_args # noqa: F401 (re-export) +from airoa_lineage.cli.dataset_utils import ( + build_input_dataset, + build_output_dataset_common_facets, +) +from airoa_lineage.facets.dataset import UsbDeviceDatasetFacet from airoa_lineage.usb_copy import USBCopySession -def _add_common_args(parser: argparse.ArgumentParser) -> None: - """ - Add common arguments shared across all commands. - - These arguments can be provided via CLI, environment variables, - or configuration file (CLI takes highest precedence). - - Args: - parser: ArgumentParser to add arguments to +def _build_input_dataset(data: Dict[str, Any], namespace: str) -> List[Dataset]: """ - parser.add_argument( - "--namespace", - help="OpenLineage namespace", - ) - parser.add_argument( - "--job-name", - help="Job name", - ) - parser.add_argument( - "--marquez-url", - help="Marquez server URL", - ) - parser.add_argument( - "--facet-prefix", - help="Facet prefix (default: empty string)", - ) - - -def _add_common_facet_args(parser: argparse.ArgumentParser) -> None: - """ - Add CommonRunFacet arguments. + Build input dataset with USB-specific facets. - These arguments represent robot and repository metadata required - for OpenLineage event tracking. + Expected data format: + { + "dataset_name": "usb_input_2025_12_08", + "usb_device": { + "id": "/dev/sdb1", + "label": "ROBOT_DATA_001", + "fs_type": "ext4" + }, + "src_dir": { + "path": "/media/usb0/robot_data/2025-12-08" + }, + "src_files": [ + {"name": "rosbag_001.bag", "hash": "abc123"}, + {"name": "rosbag_002.bag", "hash": "def456"} + ] + } Args: - parser: ArgumentParser to add arguments to - """ - parser.add_argument( - "--robot-id", - help="Robot identifier", - ) - parser.add_argument( - "--location", - help="Location identifier", - ) - parser.add_argument( - "--repository-hash", - help="Git commit hash", - ) - parser.add_argument( - "--repository-uri", - help="Repository URI", - ) - parser.add_argument( - "--repository-tag", - help="Git tag", - ) - parser.add_argument( - "--repository-branch", - help="Git branch", - ) - - -def create_parser() -> argparse.ArgumentParser: - """ - Create argument parser for airoa-lineage-usb-copy CLI. + data: Parsed JSON data + namespace: OpenLineage namespace Returns: - Configured ArgumentParser instance - """ - parser = argparse.ArgumentParser( - prog="airoa-lineage-usb-copy", - description="Track USB data copy operations with OpenLineage", - ) - - parser.add_argument( - "--version", - action="version", - version="airoa-lineage-usb-copy 0.1.0", - ) - - parser.add_argument( - "--config", - type=str, - help="Path to config file (default: ~/.config/airoa-lineage/config.json)", - ) - - parser.add_argument( - "-v", - "--verbose", - action="store_true", - help="Enable verbose output", - ) - - parser.add_argument( - "-q", - "--quiet", - action="store_true", - help="Suppress output (except errors)", - ) - - parser.add_argument( - "--json", - action="store_true", - help="Output in JSON format", - ) - - # Create subparsers for commands - subparsers = parser.add_subparsers(dest="command", required=True) - - # START command - start_parser = subparsers.add_parser( - "start", - help="Start USB copy session and emit START event", - ) - _add_common_args(start_parser) - _add_common_facet_args(start_parser) - start_parser.add_argument( - "--nominal-start-time", - help="Nominal start time (ISO 8601 format)", - ) - start_parser.add_argument( - "--nominal-end-time", - help="Nominal end time (ISO 8601 format)", - ) - start_parser.add_argument( - "--dry-run", - action="store_true", - help="Show what would be done without executing", - ) - - # COMPLETE command - complete_parser = subparsers.add_parser( - "complete", - help="Complete USB copy session and emit COMPLETE event", - ) - complete_parser.add_argument( - "--run-id", - required=True, - help="Session run ID (from start command)", - ) - _add_common_args(complete_parser) - _add_common_facet_args(complete_parser) - - # CANCEL command - cancel_parser = subparsers.add_parser( - "cancel", - help="Cancel USB copy session and emit ABORT event", - ) - cancel_parser.add_argument( - "--run-id", - required=True, - help="Session run ID (from start command)", - ) - _add_common_args(cancel_parser) - _add_common_facet_args(cancel_parser) - - return parser - - -def merge_config_with_args( - config: Dict[str, Any], - args: argparse.Namespace, -) -> Dict[str, Any]: - """ - Merge configuration with CLI arguments. + List containing single input Dataset + """ + # Build base dataset with common facets (src_dir, src_files) + dataset = build_input_dataset(data, namespace) + + # Add USB-specific facet if present + if "usb_device" in data: + usb_dev = data["usb_device"] + dataset.facets["usbDevice"] = UsbDeviceDatasetFacet( + id=usb_dev["id"], + label=usb_dev["label"], + fsType=usb_dev["fs_type"], + ) - CLI arguments take precedence over config file and environment variables. + return [dataset] - Args: - config: Configuration from load_config() - args: Parsed CLI arguments - Returns: - Merged configuration dictionary +def _build_output_dataset(data: Dict[str, Any], namespace: str) -> List[Dataset]: """ - # Top-level configuration - if args.namespace: - config["namespace"] = args.namespace - - if hasattr(args, "marquez_url") and args.marquez_url: - config["marquez_url"] = args.marquez_url - - if hasattr(args, "job_name") and args.job_name: - config["job_name"] = args.job_name - - if hasattr(args, "facet_prefix") and args.facet_prefix: - config["facet_prefix"] = args.facet_prefix - - # CommonRunFacet fields - if hasattr(args, "robot_id") and args.robot_id: - if "common_facet" not in config: - config["common_facet"] = {} - config["common_facet"]["robotId"] = args.robot_id - - if hasattr(args, "location") and args.location: - if "common_facet" not in config: - config["common_facet"] = {} - config["common_facet"]["location"] = args.location - - if hasattr(args, "repository_hash") and args.repository_hash: - if "common_facet" not in config: - config["common_facet"] = {} - config["common_facet"]["repositoryHash"] = args.repository_hash - - if hasattr(args, "repository_uri") and args.repository_uri: - if "common_facet" not in config: - config["common_facet"] = {} - config["common_facet"]["repositoryUri"] = args.repository_uri - - if hasattr(args, "repository_tag") and args.repository_tag: - if "common_facet" not in config: - config["common_facet"] = {} - config["common_facet"]["repositoryTag"] = args.repository_tag - - if hasattr(args, "repository_branch") and args.repository_branch: - if "common_facet" not in config: - config["common_facet"] = {} - config["common_facet"]["repositoryBranch"] = args.repository_branch + Build output dataset with USB-specific facets. - return config - - -def cmd_start(args: argparse.Namespace, config: Dict[str, Any]) -> int: - """ - Execute start command. + Expected data format: + { + "dataset_name": "usb_output_2025_12_08", + "dest_dir": { + "path": "/data/robot_data/2025-12-08" + }, + "operation_stats": { + "duration_seconds": 125.47 + }, + "usb_device": { # optional + "id": "/dev/sdb1", + "label": "ROBOT_DATA_001", + "fs_type": "ext4" + } + } Args: - args: Parsed CLI arguments - config: Merged configuration + data: Parsed JSON data + namespace: OpenLineage namespace Returns: - Exit code (0=success, 1=error) - """ - try: - # Build CommonRunFacet - common_facet = build_common_facet(config) - - # Validate required fields - if "namespace" not in config: - print( - "Error: namespace is required. " - "Provide via --namespace, config file, or AIROA_NAMESPACE env var.", - file=sys.stderr, - ) - return 1 - - # Dry run mode - if args.dry_run: - print("[DRY RUN] Would start session with:") - print(f" Namespace: {config['namespace']}") - print(f" Job name: {config.get('job_name', 'usb-data-copy')}") - print(f" Robot ID: {common_facet.robotId}") - print(f" Location: {common_facet.location}") - print(f" Repository: {common_facet.repositoryUri}") - print(f" Commit: {common_facet.repositoryHash}") - print(f" Branch: {common_facet.repositoryBranch}") - print(f" Tag: {common_facet.repositoryTag}") - if hasattr(args, "nominal_start_time") and args.nominal_start_time: - print(f" Nominal start: {args.nominal_start_time}") - if hasattr(args, "nominal_end_time") and args.nominal_end_time: - print(f" Nominal end: {args.nominal_end_time}") - print("[DRY RUN] Would output run_id (not generated in dry-run mode)") - return 0 - - # Create session - session = USBCopySession( - namespace=config["namespace"], - common_facet=common_facet, - job_name=config.get("job_name", "usb-data-copy"), - marquez_url=config.get("marquez_url"), - facet_prefix=config.get("facet_prefix", ""), - ) - - # Start session - run_id = session.start( - nominal_start_time=getattr(args, "nominal_start_time", None), - nominal_end_time=getattr(args, "nominal_end_time", None), - ) - - # Output run_id - if args.json: - output = {"run_id": run_id, "status": "started"} - print(json.dumps(output)) - elif args.verbose: - print("Session started successfully", file=sys.stderr) - print(run_id) - else: - # Default: just print run_id - print(run_id) - - return 0 - - except ValueError as e: - print(f"Configuration error: {e}", file=sys.stderr) - return 3 - except RuntimeError as e: - print(f"Session error: {e}", file=sys.stderr) - return 1 - except Exception as e: - print(f"Unexpected error: {e}", file=sys.stderr) - if args.verbose: - import traceback - - traceback.print_exc() - return 1 - - -def cmd_complete(args: argparse.Namespace, config: Dict[str, Any]) -> int: + List containing single output Dataset """ - Execute complete command. + # Extract dataset name + dataset_name = data.get("dataset_name", "usb_output") - Args: - args: Parsed CLI arguments - config: Merged configuration + # Get common facets + facets = build_output_dataset_common_facets(data) - Returns: - Exit code (0=success, 1=error) - """ - try: - # Build CommonRunFacet - common_facet = build_common_facet(config) - - # Validate required fields - if "namespace" not in config: - print( - "Error: namespace is required. " - "Provide via --namespace, config file, or AIROA_NAMESPACE env var.", - file=sys.stderr, - ) - return 1 - - # Create session - session = USBCopySession( - namespace=config["namespace"], - common_facet=common_facet, - job_name=config.get("job_name", "usb-data-copy"), - marquez_url=config.get("marquez_url"), - facet_prefix=config.get("facet_prefix", ""), + # Add USB-specific facet if present + if "usb_device" in data: + usb_data = data["usb_device"] + facets["usbDevice"] = UsbDeviceDatasetFacet( + id=usb_data.get("id", ""), + label=usb_data.get("label", ""), + fsType=usb_data.get("fs_type", ""), ) - # Resume existing session - session.resume(args.run_id) - - # Complete session - session.complete() + # Create dataset + dataset = Dataset(namespace=namespace, name=dataset_name, facets=facets) - # Output success message - if not args.quiet: - if args.json: - output = {"run_id": args.run_id, "status": "completed"} - print(json.dumps(output)) - else: - print(f"Session {args.run_id} completed successfully") + return [dataset] - return 0 - except ValueError as e: - print(f"Configuration error: {e}", file=sys.stderr) - return 3 - except RuntimeError as e: - print(f"Session error: {e}", file=sys.stderr) - return 1 - except Exception as e: - print(f"Unexpected error: {e}", file=sys.stderr) - if args.verbose: - import traceback +# CLI Configuration +_cli_config = CLIConfig( + prog="airoa-lineage-usb-copy", + description="Track USB data copy operations with OpenLineage", + version="airoa-lineage-usb-copy 0.1.0", + default_job_name="usb-data-copy", + session_class=USBCopySession, + build_input_dataset=_build_input_dataset, + build_output_dataset=_build_output_dataset, + start_help="Start USB copy session and emit START event", + complete_help="Complete USB copy session and emit COMPLETE event", + cancel_help="Cancel USB copy session and emit ABORT event", +) - traceback.print_exc() - return 1 - -def cmd_cancel(args: argparse.Namespace, config: Dict[str, Any]) -> int: - """ - Execute cancel command. - - Args: - args: Parsed CLI arguments - config: Merged configuration +def create_parser(): + """Create argument parser for airoa-lineage-usb-copy CLI. Returns: - Exit code (0=success, 1=error) + Configured ArgumentParser instance """ - try: - # Build CommonRunFacet - common_facet = build_common_facet(config) - - # Validate required fields - if "namespace" not in config: - print( - "Error: namespace is required. " - "Provide via --namespace, config file, or AIROA_NAMESPACE env var.", - file=sys.stderr, - ) - return 1 - - # Create session - session = USBCopySession( - namespace=config["namespace"], - common_facet=common_facet, - job_name=config.get("job_name", "usb-data-copy"), - marquez_url=config.get("marquez_url"), - facet_prefix=config.get("facet_prefix", ""), - ) - - # Resume existing session - session.resume(args.run_id) - - # Cancel session - session.cancel() - - # Output success message - if not args.quiet: - if args.json: - output = {"run_id": args.run_id, "status": "cancelled"} - print(json.dumps(output)) - else: - print(f"Session {args.run_id} cancelled successfully") - - return 0 - - except ValueError as e: - print(f"Configuration error: {e}", file=sys.stderr) - return 3 - except RuntimeError as e: - print(f"Session error: {e}", file=sys.stderr) - return 1 - except Exception as e: - print(f"Unexpected error: {e}", file=sys.stderr) - if args.verbose: - import traceback - - traceback.print_exc() - return 1 + return _create_parser(_cli_config) def main() -> int: @@ -465,36 +152,7 @@ def main() -> int: Returns: Exit code (0=success, 1=general error, 2=connection error, 3=config error) """ - parser = create_parser() - args = parser.parse_args() - - # Load configuration - try: - config_path = Path(args.config) if args.config else None - config = load_config(config_path) - except FileNotFoundError as e: - print(f"Configuration file not found: {e}", file=sys.stderr) - return 3 - except json.JSONDecodeError as e: - print(f"Invalid JSON in configuration file: {e}", file=sys.stderr) - return 3 - except Exception as e: - print(f"Error loading configuration: {e}", file=sys.stderr) - return 3 - - # Merge with CLI arguments - config = merge_config_with_args(config, args) - - # Execute command - if args.command == "start": - return cmd_start(args, config) - elif args.command == "complete": - return cmd_complete(args, config) - elif args.command == "cancel": - return cmd_cancel(args, config) - else: - print(f"Unknown command: {args.command}", file=sys.stderr) - return 1 + return run_cli(_cli_config) if __name__ == "__main__": diff --git a/src/airoa_lineage/conversion/session.py b/src/airoa_lineage/conversion/session.py index 159d81e..b78a6ae 100644 --- a/src/airoa_lineage/conversion/session.py +++ b/src/airoa_lineage/conversion/session.py @@ -258,7 +258,7 @@ def complete( # type: ignore[override] # Call parent's complete() method return super().complete() - def running(self, message: Optional[str] = None) -> None: + def running(self) -> None: """ Send RUNNING event to Marquez. @@ -266,9 +266,6 @@ def running(self, message: Optional[str] = None) -> None: of the data conversion session. Can be called multiple times between start() and complete(). - Args: - message: Optional status message describing current progress - Raises: RuntimeError: If session was not started or already completed @@ -296,10 +293,8 @@ def running(self, message: Optional[str] = None) -> None: ... aws_job_facet=aws_job_facet ... ) >>> session.start(input_datasets=[input_ds]) - >>> session.running(message="Phase 1: Reading rosbag files") - >>> session.running(message="Phase 2: Converting to LeRobot format") + >>> session.running() + >>> session.running() >>> session.complete(output_datasets=[output_ds]) """ - # Note: message parameter is kept for backward compatibility but not used - # Call parent's running() method return super().running() diff --git a/src/airoa_lineage/facets/__init__.py b/src/airoa_lineage/facets/__init__.py index fc07ab5..0d653e8 100644 --- a/src/airoa_lineage/facets/__init__.py +++ b/src/airoa_lineage/facets/__init__.py @@ -1,7 +1,29 @@ """OpenLineage facets for AIRoA lineage tracking.""" -from airoa_lineage.facets.aws_job import AWSJobRunFacet -from airoa_lineage.facets.common import CommonRunFacet -from airoa_lineage.facets.device import DeviceRunFacet +from airoa_lineage.facets.dataset import ( + DestDirDatasetFacet, + OperationStatsDatasetFacet, + SrcDirDatasetFacet, + SrcFileInfo, + SrcFilesDatasetFacet, + UsbDeviceDatasetFacet, +) +from airoa_lineage.facets.run import ( + AWSJobRunFacet, + CommonRunFacet, + DeviceRunFacet, + JobRunFacet, +) -__all__ = ["AWSJobRunFacet", "CommonRunFacet", "DeviceRunFacet"] +__all__ = [ + "AWSJobRunFacet", + "CommonRunFacet", + "DestDirDatasetFacet", + "DeviceRunFacet", + "JobRunFacet", + "OperationStatsDatasetFacet", + "SrcDirDatasetFacet", + "SrcFileInfo", + "SrcFilesDatasetFacet", + "UsbDeviceDatasetFacet", +] diff --git a/src/airoa_lineage/facets/dataset/__init__.py b/src/airoa_lineage/facets/dataset/__init__.py new file mode 100644 index 0000000..09ef2a5 --- /dev/null +++ b/src/airoa_lineage/facets/dataset/__init__.py @@ -0,0 +1,22 @@ +"""Dataset Facets for OpenLineage. + +This module provides custom OpenLineage dataset facets for tracking +metadata specific to data operations (USB, S3, NAS, etc.). +""" + +from airoa_lineage.facets.dataset.directory import ( + DestDirDatasetFacet, + SrcDirDatasetFacet, +) +from airoa_lineage.facets.dataset.files import SrcFileInfo, SrcFilesDatasetFacet +from airoa_lineage.facets.dataset.operation_stats import OperationStatsDatasetFacet +from airoa_lineage.facets.dataset.usb_device import UsbDeviceDatasetFacet + +__all__ = [ + "DestDirDatasetFacet", + "OperationStatsDatasetFacet", + "SrcDirDatasetFacet", + "SrcFileInfo", + "SrcFilesDatasetFacet", + "UsbDeviceDatasetFacet", +] diff --git a/src/airoa_lineage/facets/dataset/directory.py b/src/airoa_lineage/facets/dataset/directory.py new file mode 100644 index 0000000..0795be1 --- /dev/null +++ b/src/airoa_lineage/facets/dataset/directory.py @@ -0,0 +1,87 @@ +"""Directory Dataset Facets for OpenLineage. + +This module provides generic directory path facets for source and destination +directories. These facets can be reused across different session types +(USBCopySession, S3UploadSession, NASCopySession, etc.). +""" + +import attr +from openlineage.client.facet import BaseFacet + + +@attr.s +class SrcDirDatasetFacet(BaseFacet): + """Source directory path facet. + + This facet captures the source directory path for copy operations. + It is typically attached to input datasets. + + Attributes: + path: Source directory path (e.g., `/media/usb0/robot_data/2025-12-08`) + + Example: + ```python + from openlineage.client.run import Dataset + from airoa_lineage.facets.dataset import SrcDirDatasetFacet + + src_dir_facet = SrcDirDatasetFacet( + path="/media/usb0/robot_data/2025-12-08" + ) + + input_dataset = Dataset( + namespace="airoa_examples", + name="usb_input_2025_12_08", + facets={"srcDir": src_dir_facet} + ) + ``` + """ + + path: str = attr.ib() + + @classmethod + def get_key(cls) -> str: + """Return the facet key for OpenLineage event serialization. + + Returns: + The string "srcDir" (camelCase following OpenLineage convention) + """ + return "srcDir" + + +@attr.s +class DestDirDatasetFacet(BaseFacet): + """Destination directory path facet. + + This facet captures the destination directory path for copy operations. + It is typically attached to output datasets. + + Attributes: + path: Destination directory path (e.g., `/data/robot_data/2025-12-08`) + + Example: + ```python + from openlineage.client.run import Dataset + from airoa_lineage.facets.dataset import DestDirDatasetFacet + + dest_dir_facet = DestDirDatasetFacet( + path="/data/robot_data/2025-12-08" + ) + + output_dataset = Dataset( + namespace="airoa_examples", + name="usb_output_2025_12_08", + facets={"destDir": dest_dir_facet} + ) + ``` + """ + + path: str = attr.ib() + + @classmethod + def get_key(cls) -> str: + """Return the facet key for OpenLineage event serialization. + + Returns: + The string "destDir" (camelCase following OpenLineage convention) + """ + return "destDir" diff --git a/src/airoa_lineage/facets/dataset/files.py b/src/airoa_lineage/facets/dataset/files.py new file mode 100644 index 0000000..e4f9cbd --- /dev/null +++ b/src/airoa_lineage/facets/dataset/files.py @@ -0,0 +1,84 @@ +"""Files Dataset Facet for OpenLineage. + +This module provides a generic facet for tracking source file information +including file names and MD5 hashes. This facet can be reused across different +session types (USBCopySession, S3UploadSession, NASCopySession, etc.). +""" + +from typing import List + +import attr +from openlineage.client.facet import BaseFacet + + +@attr.s +class SrcFileInfo: + """Source file information. + + This class captures metadata for a single source file in a copy operation. + + Attributes: + name: File name (e.g., `rosbag_001.bag`) + hash: MD5 hash of the file (e.g., `5d41402abc4b2a76b9719d911017c592`) + + Example: + ```python + from airoa_lineage.facets.dataset import SrcFileInfo + + file_info = SrcFileInfo( + name="rosbag_001.bag", + hash="5d41402abc4b2a76b9719d911017c592" + ) + ``` + """ + + name: str = attr.ib() + hash: str = attr.ib() + + +@attr.s +class SrcFilesDatasetFacet(BaseFacet): + """Source files list facet. + + This facet captures a list of source files with their metadata (name and MD5 hash). + It is typically attached to input datasets in copy operations. + + Attributes: + files: List of source file information objects + + Example: + ```python + from openlineage.client.run import Dataset + from airoa_lineage.facets.dataset import SrcFilesDatasetFacet, SrcFileInfo + + files_facet = SrcFilesDatasetFacet( + files=[ + SrcFileInfo( + name="rosbag_001.bag", + hash="5d41402abc4b2a76b9719d911017c592" + ), + SrcFileInfo( + name="rosbag_002.bag", + hash="7d793037a0760186574b0282f2f435e7" + ), + ] + ) + + input_dataset = Dataset( + namespace="airoa_examples", + name="usb_input_2025_12_08", + facets={"srcFiles": files_facet} + ) + ``` + """ + + files: List[SrcFileInfo] = attr.ib() + + @classmethod + def get_key(cls) -> str: + """Return the facet key for OpenLineage event serialization. + + Returns: + The string "srcFiles" (camelCase following OpenLineage convention) + """ + return "srcFiles" diff --git a/src/airoa_lineage/facets/dataset/operation_stats.py b/src/airoa_lineage/facets/dataset/operation_stats.py new file mode 100644 index 0000000..083e492 --- /dev/null +++ b/src/airoa_lineage/facets/dataset/operation_stats.py @@ -0,0 +1,52 @@ +"""Operation Stats Dataset Facet for OpenLineage. + +This module provides a generic facet for tracking data operation statistics +such as duration. This facet can be reused across different session types +and operation types (copy, upload, download, sync, etc.). +""" + +import attr +from openlineage.client.facet import BaseFacet + + +@attr.s +class OperationStatsDatasetFacet(BaseFacet): + """Data operation statistics facet. + + This facet captures statistics about data operations (copy, upload, download, sync, etc.), + such as the duration in seconds. It is typically attached to output datasets. + + Attributes: + durationSeconds: Operation duration in seconds (supports fractional values) + + Example: + ```python + from openlineage.client.run import Dataset + from airoa_lineage.facets.dataset import OperationStatsDatasetFacet + + stats_facet = OperationStatsDatasetFacet( + durationSeconds=125.47 # 2 minutes 5.47 seconds + ) + + output_dataset = Dataset( + namespace="airoa_examples", + name="output_2025_12_08", + facets={"operationStats": stats_facet} + ) + ``` + + Note: + The duration is measured by the CLI (shell script) by recording timestamps + before and after the operation, then calculating the difference. + """ + + durationSeconds: float = attr.ib() + + @classmethod + def get_key(cls) -> str: + """Return the facet key for OpenLineage event serialization. + + Returns: + The string "operationStats" (camelCase following OpenLineage convention) + """ + return "operationStats" diff --git a/src/airoa_lineage/facets/dataset/usb_device.py b/src/airoa_lineage/facets/dataset/usb_device.py new file mode 100644 index 0000000..ee685a9 --- /dev/null +++ b/src/airoa_lineage/facets/dataset/usb_device.py @@ -0,0 +1,53 @@ +"""USB Device Dataset Facet for OpenLineage. + +This facet captures USB device information for input datasets in USB copy operations. +""" + +import attr +from openlineage.client.facet import BaseFacet + + +@attr.s +class UsbDeviceDatasetFacet(BaseFacet): + """USB device information facet. + + This facet captures USB device metadata such as filesystem ID, disk label, + and filesystem type. It is typically attached to input datasets in USB copy + operations. + + Attributes: + id: System filesystem ID (e.g., `/dev/sdb1`) + label: USB disk label (e.g., `ROBOT_DATA_001`) + fsType: Filesystem type (e.g., `ext4`, `vfat`, `exfat`) + + Example: + ```python + from openlineage.client.run import Dataset + from airoa_lineage.facets.dataset import UsbDeviceDatasetFacet + + usb_facet = UsbDeviceDatasetFacet( + id="/dev/sdb1", + label="ROBOT_DATA_001", + fsType="ext4" + ) + + input_dataset = Dataset( + namespace="airoa_examples", + name="usb_input_2025_12_08", + facets={"usbDevice": usb_facet} + ) + ``` + """ + + id: str = attr.ib() + label: str = attr.ib() + fsType: str = attr.ib() + + @classmethod + def get_key(cls) -> str: + """Return the facet key for OpenLineage event serialization. + + Returns: + The string "usbDevice" (camelCase following OpenLineage convention) + """ + return "usbDevice" diff --git a/src/airoa_lineage/facets/run/__init__.py b/src/airoa_lineage/facets/run/__init__.py new file mode 100644 index 0000000..cb9547e --- /dev/null +++ b/src/airoa_lineage/facets/run/__init__.py @@ -0,0 +1,17 @@ +"""Run Facets for OpenLineage. + +This module provides custom OpenLineage run facets for tracking +metadata specific to job executions (teleoperation, conversion, etc.). +""" + +from airoa_lineage.facets.run.aws_job import AWSJobRunFacet +from airoa_lineage.facets.run.common import CommonRunFacet +from airoa_lineage.facets.run.device import DeviceRunFacet +from airoa_lineage.facets.run.job import JobRunFacet + +__all__ = [ + "AWSJobRunFacet", + "CommonRunFacet", + "DeviceRunFacet", + "JobRunFacet", +] diff --git a/src/airoa_lineage/facets/aws_job.py b/src/airoa_lineage/facets/run/aws_job.py similarity index 100% rename from src/airoa_lineage/facets/aws_job.py rename to src/airoa_lineage/facets/run/aws_job.py diff --git a/src/airoa_lineage/facets/common.py b/src/airoa_lineage/facets/run/common.py similarity index 100% rename from src/airoa_lineage/facets/common.py rename to src/airoa_lineage/facets/run/common.py diff --git a/src/airoa_lineage/facets/device.py b/src/airoa_lineage/facets/run/device.py similarity index 100% rename from src/airoa_lineage/facets/device.py rename to src/airoa_lineage/facets/run/device.py diff --git a/src/airoa_lineage/facets/run/job.py b/src/airoa_lineage/facets/run/job.py new file mode 100644 index 0000000..94beffb --- /dev/null +++ b/src/airoa_lineage/facets/run/job.py @@ -0,0 +1,43 @@ +"""Job metadata facet for OpenLineage tracking.""" + +import attr +from openlineage.client.facet import BaseFacet + + +@attr.s +class JobRunFacet(BaseFacet): + """ + Run facet for job identification metadata. + + This facet provides a unique identifier for tracking individual job executions. + It is designed to be extensible for future fields such as job name, type, or priority. + + Attributes: + id: Unique job identifier (e.g., UUID) + + Examples: + >>> import uuid + >>> from airoa_lineage.facets import JobRunFacet + >>> + >>> # Generate a unique job ID + >>> job_id = str(uuid.uuid4()) + >>> job_facet = JobRunFacet(id=job_id) + >>> print(job_facet.id) + 'a1b2c3d4-e5f6-7890-abcd-ef1234567890' + >>> + >>> # Use with session + >>> from airoa_lineage.usb_copy import USBCopySession + >>> session = USBCopySession( + ... namespace="production", + ... common_facet=common_facet, + ... device_facet=device_facet, + ... job_facet=job_facet + ... ) + """ + + id: str = attr.ib() + + @classmethod + def get_key(cls) -> str: + """Return the facet key for use in run facets dictionary.""" + return "job" diff --git a/src/airoa_lineage/s3_upload/session.py b/src/airoa_lineage/s3_upload/session.py index 37b1e4c..a1d5729 100644 --- a/src/airoa_lineage/s3_upload/session.py +++ b/src/airoa_lineage/s3_upload/session.py @@ -1,11 +1,12 @@ """S3 data upload session management with OpenLineage tracking.""" -from typing import Dict, Optional +from typing import Dict, List, Optional from openlineage.client.facet import BaseFacet +from openlineage.client.run import Dataset from airoa_lineage.core import BaseSession -from airoa_lineage.facets import CommonRunFacet +from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet, JobRunFacet class S3UploadSession(BaseSession): @@ -14,11 +15,22 @@ class S3UploadSession(BaseSession): This class handles the lifecycle of an S3 object storage upload operation, automatically managing the run_id and ensuring proper event sequencing. + It tracks input datasets (source directory and files) and output datasets + (destination path and upload statistics). Examples: - >>> # Basic usage - >>> from airoa_lineage.facets import CommonRunFacet + >>> # Basic usage with datasets + >>> import uuid + >>> from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet, JobRunFacet + >>> from airoa_lineage.facets.dataset import ( + ... SrcDirDatasetFacet, + ... SrcFilesDatasetFacet, + ... SrcFileInfo, + ... DestDirDatasetFacet, + ... OperationStatsDatasetFacet, + ... ) >>> from airoa_lineage.s3_upload import S3UploadSession + >>> from openlineage.client.run import Dataset >>> common_facet = CommonRunFacet( ... robotId="hsr001", ... location="weblab", @@ -27,29 +39,73 @@ class S3UploadSession(BaseSession): ... repositoryTag="v1.0.0", ... repositoryBranch="main" ... ) + >>> device_facet = DeviceRunFacet(hostname="s3-upload-server-01") + >>> job_facet = JobRunFacet(id=str(uuid.uuid4())) + >>> # Create input dataset + >>> input_ds = Dataset( + ... namespace="airoa_s3_upload", + ... name="s3_input_2025_12_08", + ... facets={ + ... "srcDir": SrcDirDatasetFacet( + ... path="/data/robot_data/2025-12-08" + ... ), + ... "srcFiles": SrcFilesDatasetFacet( + ... files=[ + ... SrcFileInfo(name="rosbag_001.bag", hash="abc123"), + ... ] + ... ), + ... } + ... ) + >>> # Create output dataset + >>> output_ds = Dataset( + ... namespace="airoa_s3_upload", + ... name="s3_output_2025_12_08", + ... facets={ + ... "destDir": DestDirDatasetFacet( + ... path="s3://my-robot-data-bucket/robot_data/2025-12-08" + ... ), + ... "operationStats": OperationStatsDatasetFacet( + ... durationSeconds=325.89 + ... ), + ... } + ... ) >>> session = S3UploadSession( - ... namespace="my_namespace", - ... common_facet=common_facet + ... namespace="airoa_s3_upload", + ... common_facet=common_facet, + ... device_facet=device_facet, + ... job_facet=job_facet + ... ) + >>> run_id = session.start( + ... input_datasets=[input_ds], + ... nominal_start_time="2025-12-08T00:00:00+00:00", + ... nominal_end_time="2025-12-08T05:00:00+00:00" ... ) - >>> run_id = session.start() >>> # Upload data to S3... - >>> session.complete() + >>> session.complete(output_datasets=[output_ds]) >>> # Custom configuration >>> session = S3UploadSession( ... namespace="production", ... common_facet=common_facet, + ... device_facet=device_facet, + ... job_facet=job_facet, ... job_name="my-s3-upload-job", ... run_id="existing-run-id-123" ... ) - >>> session.start() - >>> session.complete() + >>> session.start( + ... input_datasets=[input_ds], + ... nominal_start_time="2025-12-08T00:00:00+00:00", + ... nominal_end_time="2025-12-08T05:00:00+00:00" + ... ) + >>> session.complete(output_datasets=[output_ds]) """ def __init__( self, namespace: str, common_facet: CommonRunFacet, + device_facet: DeviceRunFacet, + job_facet: JobRunFacet, job_name: str = "s3-data-upload", marquez_url: Optional[str] = None, run_id: Optional[str] = None, @@ -61,15 +117,18 @@ def __init__( Args: namespace: OpenLineage namespace (required) common_facet: Common metadata (required, includes robotId, location, repository info) + device_facet: Device metadata (required, includes hostname) + job_facet: Job metadata (required, includes unique job ID) job_name: Job name (default: "s3-data-upload") marquez_url: Marquez server URL (default: from MARQUEZ_URL env or localhost:9000) run_id: Session run ID (auto-generated UUID if not provided) facet_prefix: Prefix for custom facet names (default: "", no prefix). - When set (e.g., "airoa"), facet key becomes "{prefix}_common". + When set (e.g., "airoa"), facet keys become "{prefix}_common", "{prefix}_device", and "{prefix}_job". Examples: >>> # Basic usage - >>> from airoa_lineage.facets import CommonRunFacet + >>> import uuid + >>> from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet, JobRunFacet >>> from airoa_lineage.s3_upload import S3UploadSession >>> common_facet = CommonRunFacet( ... robotId="hsr001", @@ -79,15 +138,21 @@ def __init__( ... repositoryTag="v1.0.0", ... repositoryBranch="main" ... ) + >>> device_facet = DeviceRunFacet(hostname="s3-upload-server-01") + >>> job_facet = JobRunFacet(id=str(uuid.uuid4())) >>> session = S3UploadSession( ... namespace="my_namespace", - ... common_facet=common_facet + ... common_facet=common_facet, + ... device_facet=device_facet, + ... job_facet=job_facet ... ) >>> # Custom configuration >>> session = S3UploadSession( ... namespace="my_namespace", ... common_facet=common_facet, + ... device_facet=device_facet, + ... job_facet=job_facet, ... job_name="my_job", ... marquez_url="http://marquez.example.com:9000" ... ) @@ -96,17 +161,202 @@ def __init__( >>> session = S3UploadSession( ... namespace="my_namespace", ... common_facet=common_facet, + ... device_facet=device_facet, + ... job_facet=job_facet, ... facet_prefix="airoa" ... ) - >>> # Facet key will be "airoa_common" + >>> # Facet keys will be "airoa_common", "airoa_device", and "airoa_job" """ super().__init__(namespace, job_name, marquez_url, run_id, facet_prefix) self.common_facet = common_facet + self.device_facet = device_facet + self.job_facet = job_facet + + # Dataset tracking + self._input_datasets: List[Dataset] = [] + self._output_datasets: List[Dataset] = [] def _get_session_facets(self) -> Dict[str, BaseFacet]: """Return session-specific facets.""" - return {"common": self.common_facet} + return { + "common": self.common_facet, + "device": self.device_facet, + "job": self.job_facet, + } + + def _get_inputs(self) -> List[Dataset]: + """Return input datasets.""" + return self._input_datasets + + def _get_outputs(self) -> List[Dataset]: + """Return output datasets.""" + return self._output_datasets def _get_producer(self) -> str: """Return producer identifier.""" return "airoa-s3-system" + + def start( # type: ignore[override] + self, + input_datasets: List[Dataset], + nominal_start_time: Optional[str] = None, + nominal_end_time: Optional[str] = None, + ) -> str: + """ + Send START event to Marquez. + + This method sends an OpenLineage START event to track the beginning + of the S3 data upload session. Optionally, you can specify the nominal + time period representing the data being processed. + + Args: + input_datasets: List of input datasets (required) + nominal_start_time: Start time of the data period being processed + (ISO 8601 format, e.g., "2025-12-08T00:00:00+00:00"). + If not specified and nominal_end_time is provided, defaults to + current time. + nominal_end_time: End time of the data period being processed + (ISO 8601 format, e.g., "2025-12-08T05:00:00+00:00"). + Optional. + + Returns: + The run_id for this session + + Raises: + RuntimeError: If session was already started + + Examples: + >>> import uuid + >>> from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet, JobRunFacet + >>> from airoa_lineage.facets.dataset import ( + ... SrcDirDatasetFacet, + ... SrcFilesDatasetFacet, + ... SrcFileInfo, + ... ) + >>> from airoa_lineage.s3_upload import S3UploadSession + >>> from openlineage.client.run import Dataset + >>> common_facet = CommonRunFacet( + ... robotId="hsr001", + ... location="weblab", + ... repositoryHash="df110d5", + ... repositoryUri="https://github.com/user/repo.git", + ... repositoryTag="v1.0.0", + ... repositoryBranch="main" + ... ) + >>> device_facet = DeviceRunFacet(hostname="s3-upload-server-01") + >>> job_facet = JobRunFacet(id=str(uuid.uuid4())) + >>> # Create input dataset + >>> input_ds = Dataset( + ... namespace="airoa_s3_upload", + ... name="s3_input_2025_12_08", + ... facets={ + ... "srcDir": SrcDirDatasetFacet( + ... path="/data/robot_data/2025-12-08" + ... ), + ... "srcFiles": SrcFilesDatasetFacet( + ... files=[ + ... SrcFileInfo(name="rosbag_001.bag", hash="abc123"), + ... ] + ... ), + ... } + ... ) + >>> session = S3UploadSession( + ... namespace="airoa_s3_upload", + ... common_facet=common_facet, + ... device_facet=device_facet, + ... job_facet=job_facet + ... ) + >>> run_id = session.start( + ... input_datasets=[input_ds], + ... nominal_start_time="2025-12-08T00:00:00+00:00", + ... nominal_end_time="2025-12-08T05:00:00+00:00" + ... ) + """ + # Store input datasets + self._input_datasets = input_datasets + # Call parent's start() method + return super().start(nominal_start_time, nominal_end_time) + + def complete( # type: ignore[override] + self, + output_datasets: List[Dataset], + ) -> None: + """ + Send COMPLETE event to Marquez. + + This method sends an OpenLineage COMPLETE event to track the successful + completion of the S3 data upload session. + + Args: + output_datasets: List of output datasets (required) + + Raises: + RuntimeError: If session was not started or already completed + + Examples: + >>> import uuid + >>> from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet, JobRunFacet + >>> from airoa_lineage.facets.dataset import ( + ... SrcDirDatasetFacet, + ... SrcFilesDatasetFacet, + ... SrcFileInfo, + ... DestDirDatasetFacet, + ... OperationStatsDatasetFacet, + ... ) + >>> from airoa_lineage.s3_upload import S3UploadSession + >>> from openlineage.client.run import Dataset + >>> common_facet = CommonRunFacet( + ... robotId="hsr001", + ... location="weblab", + ... repositoryHash="df110d5", + ... repositoryUri="https://github.com/user/repo.git", + ... repositoryTag="v1.0.0", + ... repositoryBranch="main" + ... ) + >>> device_facet = DeviceRunFacet(hostname="s3-upload-server-01") + >>> job_facet = JobRunFacet(id=str(uuid.uuid4())) + >>> # Create input dataset + >>> input_ds = Dataset( + ... namespace="airoa_s3_upload", + ... name="s3_input_2025_12_08", + ... facets={ + ... "srcDir": SrcDirDatasetFacet( + ... path="/data/robot_data/2025-12-08" + ... ), + ... "srcFiles": SrcFilesDatasetFacet( + ... files=[ + ... SrcFileInfo(name="rosbag_001.bag", hash="abc123"), + ... ] + ... ), + ... } + ... ) + >>> # Create output dataset + >>> output_ds = Dataset( + ... namespace="airoa_s3_upload", + ... name="s3_output_2025_12_08", + ... facets={ + ... "destDir": DestDirDatasetFacet( + ... path="s3://my-robot-data-bucket/robot_data/2025-12-08" + ... ), + ... "operationStats": OperationStatsDatasetFacet( + ... durationSeconds=325.89 + ... ), + ... } + ... ) + >>> session = S3UploadSession( + ... namespace="airoa_s3_upload", + ... common_facet=common_facet, + ... device_facet=device_facet, + ... job_facet=job_facet + ... ) + >>> session.start( + ... input_datasets=[input_ds], + ... nominal_start_time="2025-12-08T00:00:00+00:00", + ... nominal_end_time="2025-12-08T05:00:00+00:00" + ... ) + >>> session.complete(output_datasets=[output_ds]) + """ + # Store output datasets + self._output_datasets = output_datasets + # Call parent's complete() method + return super().complete() diff --git a/src/airoa_lineage/usb_copy/session.py b/src/airoa_lineage/usb_copy/session.py index cfb3119..d8c20f5 100644 --- a/src/airoa_lineage/usb_copy/session.py +++ b/src/airoa_lineage/usb_copy/session.py @@ -1,11 +1,12 @@ """USB data copy session management with OpenLineage tracking.""" -from typing import Dict, Optional +from typing import Dict, List, Optional from openlineage.client.facet import BaseFacet +from openlineage.client.run import Dataset from airoa_lineage.core import BaseSession -from airoa_lineage.facets import CommonRunFacet +from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet, JobRunFacet class USBCopySession(BaseSession): @@ -16,9 +17,19 @@ class USBCopySession(BaseSession): managing the run_id and ensuring proper event sequencing. Examples: - >>> # Basic usage - >>> from airoa_lineage.facets import CommonRunFacet + >>> # Basic usage with datasets + >>> import uuid + >>> from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet, JobRunFacet + >>> from airoa_lineage.facets.dataset import ( + ... UsbDeviceDatasetFacet, + ... SrcDirDatasetFacet, + ... SrcFilesDatasetFacet, + ... SrcFileInfo, + ... DestDirDatasetFacet, + ... OperationStatsDatasetFacet, + ... ) >>> from airoa_lineage.usb_copy import USBCopySession + >>> from openlineage.client.run import Dataset >>> common_facet = CommonRunFacet( ... robotId="hsr001", ... location="weblab", @@ -27,29 +38,59 @@ class USBCopySession(BaseSession): ... repositoryTag="v1.0.0", ... repositoryBranch="main" ... ) - >>> session = USBCopySession( - ... namespace="my_namespace", - ... common_facet=common_facet + >>> device_facet = DeviceRunFacet(hostname="copy-pc-001") + >>> job_facet = JobRunFacet(id=str(uuid.uuid4())) + >>> # Create input dataset with USB device, source directory, and files info + >>> input_ds = Dataset( + ... namespace="airoa_usb_copy", + ... name="usb_input_2025_12_08", + ... facets={ + ... "usbDevice": UsbDeviceDatasetFacet( + ... id="/dev/sdb1", + ... label="ROBOT_DATA_001", + ... fsType="ext4" + ... ), + ... "srcDir": SrcDirDatasetFacet( + ... path="/media/usb0/robot_data/2025-12-08" + ... ), + ... "srcFiles": SrcFilesDatasetFacet( + ... files=[ + ... SrcFileInfo(name="rosbag_001.bag", hash="abc123"), + ... SrcFileInfo(name="rosbag_002.bag", hash="def456"), + ... ] + ... ), + ... } + ... ) + >>> # Create output dataset with destination directory and copy stats + >>> output_ds = Dataset( + ... namespace="airoa_usb_copy", + ... name="usb_output_2025_12_08", + ... facets={ + ... "destDir": DestDirDatasetFacet( + ... path="/data/robot_data/2025-12-08" + ... ), + ... "operationStats": OperationStatsDatasetFacet( + ... durationSeconds=125.47 + ... ), + ... } ... ) - >>> run_id = session.start() - >>> # Copy data from USB... - >>> session.complete() - - >>> # Custom configuration >>> session = USBCopySession( - ... namespace="production", + ... namespace="airoa_usb_copy", ... common_facet=common_facet, - ... job_name="my-usb-copy-job", - ... run_id="existing-run-id-123" + ... device_facet=device_facet, + ... job_facet=job_facet ... ) - >>> session.start() - >>> session.complete() + >>> run_id = session.start(input_datasets=[input_ds]) + >>> # Copy data from USB... + >>> session.complete(output_datasets=[output_ds]) """ def __init__( self, namespace: str, common_facet: CommonRunFacet, + device_facet: DeviceRunFacet, + job_facet: JobRunFacet, job_name: str = "usb-data-copy", marquez_url: Optional[str] = None, run_id: Optional[str] = None, @@ -61,15 +102,18 @@ def __init__( Args: namespace: OpenLineage namespace (required) common_facet: Common metadata (required, includes robotId, location, repository info) + device_facet: Device metadata (required, includes hostname) + job_facet: Job metadata (required, includes unique job ID) job_name: Job name (default: "usb-data-copy") marquez_url: Marquez server URL (default: from MARQUEZ_URL env or localhost:9000) run_id: Session run ID (auto-generated UUID if not provided) facet_prefix: Prefix for custom facet names (default: "", no prefix). - When set (e.g., "airoa"), facet key becomes "{prefix}_common". + When set (e.g., "airoa"), facet keys become "{prefix}_common", "{prefix}_device", and "{prefix}_job". Examples: >>> # Basic usage - >>> from airoa_lineage.facets import CommonRunFacet + >>> import uuid + >>> from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet, JobRunFacet >>> from airoa_lineage.usb_copy import USBCopySession >>> common_facet = CommonRunFacet( ... robotId="hsr001", @@ -79,15 +123,21 @@ def __init__( ... repositoryTag="v1.0.0", ... repositoryBranch="main" ... ) + >>> device_facet = DeviceRunFacet(hostname="copy-pc-001") + >>> job_facet = JobRunFacet(id=str(uuid.uuid4())) >>> session = USBCopySession( ... namespace="my_namespace", - ... common_facet=common_facet + ... common_facet=common_facet, + ... device_facet=device_facet, + ... job_facet=job_facet ... ) >>> # Custom configuration >>> session = USBCopySession( ... namespace="my_namespace", ... common_facet=common_facet, + ... device_facet=device_facet, + ... job_facet=job_facet, ... job_name="my_job", ... marquez_url="http://marquez.example.com:9000" ... ) @@ -96,17 +146,220 @@ def __init__( >>> session = USBCopySession( ... namespace="my_namespace", ... common_facet=common_facet, + ... device_facet=device_facet, + ... job_facet=job_facet, ... facet_prefix="airoa" ... ) - >>> # Facet key will be "airoa_common" + >>> # Facet keys will be "airoa_common", "airoa_device", and "airoa_job" """ super().__init__(namespace, job_name, marquez_url, run_id, facet_prefix) self.common_facet = common_facet + self.device_facet = device_facet + self.job_facet = job_facet + self._input_datasets: List[Dataset] = [] + self._output_datasets: List[Dataset] = [] def _get_session_facets(self) -> Dict[str, BaseFacet]: """Return session-specific facets.""" - return {"common": self.common_facet} + return { + "common": self.common_facet, + "device": self.device_facet, + "job": self.job_facet, + } def _get_producer(self) -> str: """Return producer identifier.""" return "airoa-usbcopy-system" + + def _get_inputs(self) -> List[Dataset]: + """Return input datasets.""" + return self._input_datasets + + def _get_outputs(self) -> List[Dataset]: + """Return output datasets.""" + return self._output_datasets + + def start( # type: ignore[override] + self, + input_datasets: List[Dataset], + nominal_start_time: Optional[str] = None, + nominal_end_time: Optional[str] = None, + ) -> str: + """ + Send START event to Marquez. + + This method sends an OpenLineage START event to track the beginning + of the USB data copy session. Optionally, you can specify the nominal + time period representing the data being processed. + + Args: + input_datasets: List of input datasets (required) + nominal_start_time: Start time of the data period being processed + (ISO 8601 format, e.g., "2025-12-08T00:00:00+00:00"). + If not specified and nominal_end_time is provided, defaults to + current time. + nominal_end_time: End time of the data period being processed + (ISO 8601 format, e.g., "2025-12-08T05:00:00+00:00"). + Optional. + + Returns: + The run_id for this session + + Raises: + RuntimeError: If session was already started + + Examples: + >>> # Basic usage without nominal time + >>> import uuid + >>> from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet, JobRunFacet + >>> from airoa_lineage.facets.dataset import ( + ... UsbDeviceDatasetFacet, + ... SrcDirDatasetFacet, + ... SrcFilesDatasetFacet, + ... SrcFileInfo, + ... ) + >>> from airoa_lineage.usb_copy import USBCopySession + >>> from openlineage.client.run import Dataset + >>> common_facet = CommonRunFacet( + ... robotId="hsr001", + ... location="weblab", + ... repositoryHash="df110d5", + ... repositoryUri="https://github.com/user/repo.git", + ... repositoryTag="v1.0.0", + ... repositoryBranch="main" + ... ) + >>> device_facet = DeviceRunFacet(hostname="copy-pc-001") + >>> job_facet = JobRunFacet(id=str(uuid.uuid4())) + >>> # Create input dataset with facets + >>> usb_facet = UsbDeviceDatasetFacet( + ... id="/dev/sdb1", + ... label="ROBOT_DATA_001", + ... fsType="ext4" + ... ) + >>> src_dir_facet = SrcDirDatasetFacet( + ... path="/media/usb0/robot_data/2025-12-08" + ... ) + >>> src_files_facet = SrcFilesDatasetFacet( + ... files=[ + ... SrcFileInfo(name="rosbag_001.bag", hash="abc123"), + ... SrcFileInfo(name="rosbag_002.bag", hash="def456"), + ... ] + ... ) + >>> input_ds = Dataset( + ... namespace="airoa_usb_copy", + ... name="usb_input_2025_12_08", + ... facets={ + ... "usbDevice": usb_facet, + ... "srcDir": src_dir_facet, + ... "srcFiles": src_files_facet, + ... } + ... ) + >>> session = USBCopySession( + ... namespace="airoa_usb_copy", + ... common_facet=common_facet, + ... device_facet=device_facet, + ... job_facet=job_facet + ... ) + >>> run_id = session.start(input_datasets=[input_ds]) + >>> print(f"Session started: {run_id}") + + >>> # With nominal time period + >>> run_id = session.start( + ... input_datasets=[input_ds], + ... nominal_start_time="2025-12-08T00:00:00+00:00", + ... nominal_end_time="2025-12-08T05:00:00+00:00" + ... ) + """ + # Store input datasets + self._input_datasets = input_datasets + # Call parent's start() method + return super().start(nominal_start_time, nominal_end_time) + + def complete( # type: ignore[override] + self, + output_datasets: List[Dataset], + ) -> None: + """ + Send COMPLETE event to Marquez. + + This method sends an OpenLineage COMPLETE event to track the successful + completion of the USB data copy session. + + Args: + output_datasets: List of output datasets (required) + + Raises: + RuntimeError: If session was not started or already completed + + Examples: + >>> import uuid + >>> from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet, JobRunFacet + >>> from airoa_lineage.facets.dataset import ( + ... UsbDeviceDatasetFacet, + ... SrcDirDatasetFacet, + ... SrcFilesDatasetFacet, + ... SrcFileInfo, + ... DestDirDatasetFacet, + ... OperationStatsDatasetFacet, + ... ) + >>> from airoa_lineage.usb_copy import USBCopySession + >>> from openlineage.client.run import Dataset + >>> common_facet = CommonRunFacet( + ... robotId="hsr001", + ... location="weblab", + ... repositoryHash="df110d5", + ... repositoryUri="https://github.com/user/repo.git", + ... repositoryTag="v1.0.0", + ... repositoryBranch="main" + ... ) + >>> device_facet = DeviceRunFacet(hostname="copy-pc-001") + >>> job_facet = JobRunFacet(id=str(uuid.uuid4())) + >>> # Create input dataset + >>> input_ds = Dataset( + ... namespace="airoa_usb_copy", + ... name="usb_input_2025_12_08", + ... facets={ + ... "usbDevice": UsbDeviceDatasetFacet( + ... id="/dev/sdb1", + ... label="ROBOT_DATA_001", + ... fsType="ext4" + ... ), + ... "srcDir": SrcDirDatasetFacet( + ... path="/media/usb0/robot_data/2025-12-08" + ... ), + ... "srcFiles": SrcFilesDatasetFacet( + ... files=[ + ... SrcFileInfo(name="rosbag_001.bag", hash="abc123"), + ... ] + ... ), + ... } + ... ) + >>> # Create output dataset + >>> dest_dir_facet = DestDirDatasetFacet( + ... path="/data/robot_data/2025-12-08" + ... ) + >>> operation_stats_facet = OperationStatsDatasetFacet( + ... durationSeconds=125.47 + ... ) + >>> output_ds = Dataset( + ... namespace="airoa_usb_copy", + ... name="usb_output_2025_12_08", + ... facets={ + ... "destDir": dest_dir_facet, + ... "operationStats": operation_stats_facet, + ... } + ... ) + >>> session = USBCopySession( + ... namespace="airoa_usb_copy", + ... common_facet=common_facet, + ... device_facet=device_facet, + ... job_facet=job_facet + ... ) + >>> session.start(input_datasets=[input_ds]) + >>> # Perform USB copy... + >>> session.complete(output_datasets=[output_ds]) + """ + # Store output datasets + self._output_datasets = output_datasets + # Call parent's complete() method + return super().complete() diff --git a/tests/conftest.py b/tests/conftest.py index 13d2b7d..8758ef3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,8 +1,15 @@ """Shared pytest fixtures for all tests.""" +import uuid + import pytest -from airoa_lineage.facets import AWSJobRunFacet, CommonRunFacet, DeviceRunFacet +from airoa_lineage.facets import ( + AWSJobRunFacet, + CommonRunFacet, + DeviceRunFacet, + JobRunFacet, +) @pytest.fixture @@ -49,7 +56,13 @@ def alternative_common_facet(): @pytest.fixture def device_facet(): """Standard DeviceRunFacet for testing.""" - return DeviceRunFacet(hostname="operator-pc-001") + return DeviceRunFacet(hostname="copy-pc-001") + + +@pytest.fixture +def job_facet(): + """Standard JobRunFacet for testing.""" + return JobRunFacet(id=str(uuid.uuid4())) @pytest.fixture diff --git a/tests/unit/cli/test_base.py b/tests/unit/cli/test_base.py new file mode 100644 index 0000000..978a7c8 --- /dev/null +++ b/tests/unit/cli/test_base.py @@ -0,0 +1,278 @@ +"""Tests for CLI base module.""" + +import json +import tempfile +from typing import Any, Dict, List + +import pytest +from openlineage.client.run import Dataset + +from airoa_lineage.cli.base import ( + CLIConfig, + create_parser, + load_dataset_from_file, +) +from airoa_lineage.core import BaseSession + + +# Mock session class for testing +class MockSession(BaseSession): + """Mock session class for testing.""" + + def __init__(self, **kwargs): + # Store kwargs for verification + self._init_kwargs = kwargs + + def _get_session_facets(self): + return {} + + def _get_producer(self): + return "test-producer" + + def _get_inputs(self): + return [] + + def _get_outputs(self): + return [] + + def start( + self, input_datasets=None, nominal_start_time=None, nominal_end_time=None + ): + return "test-run-id" + + def resume(self, run_id): + pass + + def complete(self, output_datasets=None): + pass + + def cancel(self): + pass + + +def mock_build_input(data: Dict[str, Any], namespace: str) -> List[Dataset]: + """Mock input dataset builder.""" + return [Dataset(namespace=namespace, name=data.get("dataset_name", "test_input"))] + + +def mock_build_output(data: Dict[str, Any], namespace: str) -> List[Dataset]: + """Mock output dataset builder.""" + return [Dataset(namespace=namespace, name=data.get("dataset_name", "test_output"))] + + +@pytest.fixture +def cli_config(): + """Create a test CLI configuration.""" + return CLIConfig( + prog="test-cli", + description="Test CLI description", + version="test-cli 0.1.0", + default_job_name="test-job", + session_class=MockSession, + build_input_dataset=mock_build_input, + build_output_dataset=mock_build_output, + start_help="Start test session", + complete_help="Complete test session", + cancel_help="Cancel test session", + ) + + +class TestCLIConfig: + """Tests for CLIConfig dataclass.""" + + def test_cli_config_creation(self, cli_config): + """Test that CLIConfig can be created with all fields.""" + assert cli_config.prog == "test-cli" + assert cli_config.description == "Test CLI description" + assert cli_config.version == "test-cli 0.1.0" + assert cli_config.default_job_name == "test-job" + assert cli_config.session_class == MockSession + assert cli_config.build_input_dataset == mock_build_input + assert cli_config.build_output_dataset == mock_build_output + + def test_cli_config_default_help_texts(self): + """Test that CLIConfig has default help texts.""" + config = CLIConfig( + prog="test", + description="Test", + version="1.0", + default_job_name="job", + session_class=MockSession, + build_input_dataset=mock_build_input, + build_output_dataset=mock_build_output, + ) + assert config.start_help == "Start session and emit START event" + assert config.complete_help == "Complete session and emit COMPLETE event" + assert config.cancel_help == "Cancel session and emit ABORT event" + + def test_cli_config_is_frozen(self, cli_config): + """Test that CLIConfig is immutable.""" + with pytest.raises(AttributeError): + cli_config.prog = "modified" + + +class TestCreateParser: + """Tests for create_parser function.""" + + def test_create_parser_basic(self, cli_config): + """Test that create_parser creates a parser with correct program name.""" + parser = create_parser(cli_config) + assert parser.prog == "test-cli" + + def test_create_parser_has_version(self, cli_config): + """Test that parser has version argument.""" + parser = create_parser(cli_config) + # Version action exits, so we just check it doesn't raise + with pytest.raises(SystemExit): + parser.parse_args(["--version"]) + + def test_create_parser_has_start_subcommand(self, cli_config): + """Test that parser has start subcommand.""" + parser = create_parser(cli_config) + args = parser.parse_args(["start"]) + assert args.command == "start" + + def test_create_parser_has_complete_subcommand(self, cli_config): + """Test that parser has complete subcommand with required run-id.""" + parser = create_parser(cli_config) + args = parser.parse_args(["complete", "--run-id", "test-id"]) + assert args.command == "complete" + assert args.run_id == "test-id" + + def test_create_parser_has_cancel_subcommand(self, cli_config): + """Test that parser has cancel subcommand with required run-id.""" + parser = create_parser(cli_config) + args = parser.parse_args(["cancel", "--run-id", "test-id"]) + assert args.command == "cancel" + assert args.run_id == "test-id" + + def test_start_has_common_args(self, cli_config): + """Test that start command has common arguments.""" + parser = create_parser(cli_config) + args = parser.parse_args( + [ + "start", + "--namespace", + "test-ns", + "--job-name", + "test-job", + "--marquez-url", + "http://localhost:9000", + ] + ) + assert args.namespace == "test-ns" + assert args.job_name == "test-job" + assert args.marquez_url == "http://localhost:9000" + + def test_start_has_facet_args(self, cli_config): + """Test that start command has facet arguments.""" + parser = create_parser(cli_config) + args = parser.parse_args( + [ + "start", + "--robot-id", + "hsr001", + "--location", + "weblab", + "--hostname", + "test-host", + "--job-id", + "test-job-id", + ] + ) + assert args.robot_id == "hsr001" + assert args.location == "weblab" + assert args.hostname == "test-host" + assert args.job_id == "test-job-id" + + def test_start_has_nominal_time_args(self, cli_config): + """Test that start command has nominal time arguments.""" + parser = create_parser(cli_config) + args = parser.parse_args( + [ + "start", + "--nominal-start-time", + "2025-01-01T00:00:00Z", + "--nominal-end-time", + "2025-01-01T01:00:00Z", + ] + ) + assert args.nominal_start_time == "2025-01-01T00:00:00Z" + assert args.nominal_end_time == "2025-01-01T01:00:00Z" + + def test_start_has_input_dataset_args(self, cli_config): + """Test that start command has input dataset arguments.""" + parser = create_parser(cli_config) + args = parser.parse_args(["start", "--input-dataset", '{"name": "test"}']) + assert args.input_dataset == '{"name": "test"}' + + def test_start_has_dry_run(self, cli_config): + """Test that start command has dry-run argument.""" + parser = create_parser(cli_config) + args = parser.parse_args(["start", "--dry-run"]) + assert args.dry_run is True + + def test_complete_has_output_dataset_args(self, cli_config): + """Test that complete command has output dataset arguments.""" + parser = create_parser(cli_config) + args = parser.parse_args( + ["complete", "--run-id", "test", "--output-dataset", '{"name": "test"}'] + ) + assert args.output_dataset == '{"name": "test"}' + + +class TestLoadDatasetFromFile: + """Tests for load_dataset_from_file function.""" + + def test_load_dataset_from_valid_file(self): + """Test loading dataset from a valid JSON file.""" + data = {"dataset_name": "test_dataset", "key": "value"} + + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + json.dump(data, f) + temp_path = f.name + + try: + result = load_dataset_from_file( + temp_path, "test-namespace", mock_build_input, "input" + ) + assert len(result) == 1 + assert result[0].name == "test_dataset" + assert result[0].namespace == "test-namespace" + finally: + import os + + os.unlink(temp_path) + + def test_load_dataset_file_not_found(self): + """Test loading dataset from non-existent file.""" + with pytest.raises(ValueError) as exc_info: + load_dataset_from_file( + "/nonexistent/path.json", "test-namespace", mock_build_input, "input" + ) + assert "Input dataset file not found" in str(exc_info.value) + + def test_load_dataset_invalid_json(self): + """Test loading dataset from file with invalid JSON.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + f.write("not valid json") + temp_path = f.name + + try: + with pytest.raises(ValueError) as exc_info: + load_dataset_from_file( + temp_path, "test-namespace", mock_build_input, "input" + ) + assert "Invalid JSON in input dataset file" in str(exc_info.value) + finally: + import os + + os.unlink(temp_path) + + def test_load_dataset_output_type_message(self): + """Test that output dataset type appears in error message.""" + with pytest.raises(ValueError) as exc_info: + load_dataset_from_file( + "/nonexistent/path.json", "test-namespace", mock_build_output, "output" + ) + assert "Output dataset file not found" in str(exc_info.value) diff --git a/tests/unit/cli/test_common_args.py b/tests/unit/cli/test_common_args.py new file mode 100644 index 0000000..7ec9a6e --- /dev/null +++ b/tests/unit/cli/test_common_args.py @@ -0,0 +1,386 @@ +"""Tests for CLI common_args module.""" + +import argparse + +from airoa_lineage.cli.common_args import ( + add_common_args, + add_common_facet_args, + add_device_facet_args, + add_job_facet_args, + add_repository_args, + merge_config_with_args, +) + + +class TestAddCommonArgs: + """Test add_common_args function.""" + + def test_add_common_args(self): + """Test that add_common_args adds expected arguments.""" + parser = argparse.ArgumentParser() + add_common_args(parser) + + # Parse with arguments + args = parser.parse_args( + [ + "--namespace", + "test", + "--job-name", + "test-job", + "--marquez-url", + "http://localhost:9000", + "--facet-prefix", + "test-prefix", + ] + ) + + assert args.namespace == "test" + assert args.job_name == "test-job" + assert args.marquez_url == "http://localhost:9000" + assert args.facet_prefix == "test-prefix" + + def test_add_common_args_with_defaults(self): + """Test that add_common_args handles default values.""" + parser = argparse.ArgumentParser() + add_common_args(parser) + + # Parse with minimal arguments (facet_prefix has default) + args = parser.parse_args([]) + + assert args.facet_prefix == "" + assert args.namespace is None + assert args.job_name is None + assert args.marquez_url is None + + +class TestAddCommonFacetArgs: + """Test add_common_facet_args function.""" + + def test_add_common_facet_args(self): + """Test that add_common_facet_args adds expected arguments.""" + parser = argparse.ArgumentParser() + add_common_facet_args(parser) + + # Parse with arguments + args = parser.parse_args( + [ + "--robot-id", + "hsr001", + "--location", + "weblab", + ] + ) + + assert args.robot_id == "hsr001" + assert args.location == "weblab" + + def test_add_common_facet_args_optional(self): + """Test that add_common_facet_args arguments are optional.""" + parser = argparse.ArgumentParser() + add_common_facet_args(parser) + + # Parse without arguments + args = parser.parse_args([]) + + assert args.robot_id is None + assert args.location is None + + +class TestAddRepositoryArgs: + """Test add_repository_args function.""" + + def test_add_repository_args(self): + """Test that add_repository_args adds expected arguments.""" + parser = argparse.ArgumentParser() + add_repository_args(parser) + + # Parse with arguments + args = parser.parse_args( + [ + "--repository-hash", + "abc123", + "--repository-uri", + "https://github.com/test/repo.git", + "--repository-tag", + "v1.0.0", + "--repository-branch", + "main", + ] + ) + + assert args.repository_hash == "abc123" + assert args.repository_uri == "https://github.com/test/repo.git" + assert args.repository_tag == "v1.0.0" + assert args.repository_branch == "main" + + def test_add_repository_args_optional(self): + """Test that add_repository_args arguments are optional.""" + parser = argparse.ArgumentParser() + add_repository_args(parser) + + # Parse without arguments + args = parser.parse_args([]) + + assert args.repository_hash is None + assert args.repository_uri is None + assert args.repository_tag is None + assert args.repository_branch is None + + +class TestAddDeviceFacetArgs: + """Test add_device_facet_args function.""" + + def test_add_device_facet_args(self): + """Test that add_device_facet_args adds expected arguments.""" + parser = argparse.ArgumentParser() + add_device_facet_args(parser) + + # Parse with arguments + args = parser.parse_args( + [ + "--hostname", + "test-host", + ] + ) + + assert args.hostname == "test-host" + + def test_add_device_facet_args_optional(self): + """Test that add_device_facet_args argument is optional.""" + parser = argparse.ArgumentParser() + add_device_facet_args(parser) + + # Parse without arguments + args = parser.parse_args([]) + + assert args.hostname is None + + +class TestAddJobFacetArgs: + """Test add_job_facet_args function.""" + + def test_add_job_facet_args(self): + """Test that add_job_facet_args adds expected arguments.""" + parser = argparse.ArgumentParser() + add_job_facet_args(parser) + + # Parse with arguments + args = parser.parse_args( + [ + "--job-id", + "550e8400-e29b-41d4-a716-446655440000", + ] + ) + + assert args.job_id == "550e8400-e29b-41d4-a716-446655440000" + + def test_add_job_facet_args_optional(self): + """Test that add_job_facet_args argument is optional.""" + parser = argparse.ArgumentParser() + add_job_facet_args(parser) + + # Parse without arguments + args = parser.parse_args([]) + + assert args.job_id is None + + +class TestCombinedArgs: + """Test combining multiple argument functions.""" + + def test_all_args_combined(self): + """Test that all argument functions can be combined.""" + parser = argparse.ArgumentParser() + add_common_args(parser) + add_common_facet_args(parser) + add_repository_args(parser) + add_device_facet_args(parser) + add_job_facet_args(parser) + + # Parse with all arguments + args = parser.parse_args( + [ + "--namespace", + "test", + "--job-name", + "test-job", + "--marquez-url", + "http://localhost:9000", + "--facet-prefix", + "test-prefix", + "--robot-id", + "hsr001", + "--location", + "weblab", + "--repository-hash", + "abc123", + "--repository-uri", + "https://github.com/test/repo.git", + "--repository-tag", + "v1.0.0", + "--repository-branch", + "main", + "--hostname", + "test-host", + "--job-id", + "550e8400-e29b-41d4-a716-446655440000", + ] + ) + + # Verify all arguments are present + assert args.namespace == "test" + assert args.job_name == "test-job" + assert args.marquez_url == "http://localhost:9000" + assert args.facet_prefix == "test-prefix" + assert args.robot_id == "hsr001" + assert args.location == "weblab" + assert args.repository_hash == "abc123" + assert args.repository_uri == "https://github.com/test/repo.git" + assert args.repository_tag == "v1.0.0" + assert args.repository_branch == "main" + assert args.hostname == "test-host" + assert args.job_id == "550e8400-e29b-41d4-a716-446655440000" + + +class TestMergeConfigWithArgs: + """Test merge_config_with_args function.""" + + def test_top_level_args_override_config(self): + """Test that CLI args override top-level config values.""" + config = { + "namespace": "original", + "marquez_url": "http://original:9000", + "job_name": "original-job", + "facet_prefix": "original-prefix", + } + args = argparse.Namespace( + namespace="cli-namespace", + marquez_url="http://cli:9000", + job_name="cli-job", + facet_prefix="cli-prefix", + ) + + result = merge_config_with_args(config, args) + + assert result["namespace"] == "cli-namespace" + assert result["marquez_url"] == "http://cli:9000" + assert result["job_name"] == "cli-job" + assert result["facet_prefix"] == "cli-prefix" + + def test_nested_args_create_sections(self): + """Test that nested args create config sections if not present.""" + config = {} + args = argparse.Namespace( + namespace=None, + marquez_url=None, + job_name=None, + facet_prefix=None, + robot_id="hsr001", + location="weblab", + hostname="test-host", + job_id="test-job-id", + ) + + result = merge_config_with_args(config, args) + + assert result["common_facet"]["robotId"] == "hsr001" + assert result["common_facet"]["location"] == "weblab" + assert result["device_facet"]["hostname"] == "test-host" + assert result["job_facet"]["id"] == "test-job-id" + + def test_nested_args_preserve_existing_sections(self): + """Test that nested args preserve existing section values.""" + config = { + "common_facet": {"robotId": "original", "extra": "value"}, + } + args = argparse.Namespace( + namespace=None, + marquez_url=None, + job_name=None, + facet_prefix=None, + robot_id="hsr001", + location="weblab", + ) + + result = merge_config_with_args(config, args) + + assert result["common_facet"]["robotId"] == "hsr001" + assert result["common_facet"]["location"] == "weblab" + assert result["common_facet"]["extra"] == "value" + + def test_none_args_do_not_override(self): + """Test that None/falsy args don't override config values.""" + config = { + "namespace": "original", + "common_facet": {"robotId": "original-robot"}, + } + args = argparse.Namespace( + namespace=None, + marquez_url=None, + job_name=None, + facet_prefix=None, + robot_id=None, + location=None, + ) + + result = merge_config_with_args(config, args) + + assert result["namespace"] == "original" + assert result["common_facet"]["robotId"] == "original-robot" + + def test_repository_args(self): + """Test repository-related arguments.""" + config = {} + args = argparse.Namespace( + namespace=None, + marquez_url=None, + job_name=None, + facet_prefix=None, + robot_id=None, + location=None, + repository_hash="abc123", + repository_uri="https://github.com/test/repo.git", + repository_tag="v1.0.0", + repository_branch="main", + hostname=None, + job_id=None, + ) + + result = merge_config_with_args(config, args) + + assert result["common_facet"]["repositoryHash"] == "abc123" + assert ( + result["common_facet"]["repositoryUri"] + == "https://github.com/test/repo.git" + ) + assert result["common_facet"]["repositoryTag"] == "v1.0.0" + assert result["common_facet"]["repositoryBranch"] == "main" + + def test_missing_args_attributes(self): + """Test handling of args without all attributes (hasattr behavior).""" + config = {"namespace": "original"} + # Minimal args without most attributes + args = argparse.Namespace(namespace="cli-namespace") + + result = merge_config_with_args(config, args) + + assert result["namespace"] == "cli-namespace" + # Should not have created any nested sections + assert "common_facet" not in result + assert "device_facet" not in result + assert "job_facet" not in result + + def test_empty_string_args_do_not_override(self): + """Test that empty string args don't override config values.""" + config = {"namespace": "original"} + args = argparse.Namespace( + namespace="", + marquez_url="", + job_name="", + facet_prefix="", + ) + + result = merge_config_with_args(config, args) + + # Empty string is falsy, so original value should be preserved + assert result["namespace"] == "original" diff --git a/tests/unit/cli/test_dataset_utils.py b/tests/unit/cli/test_dataset_utils.py new file mode 100644 index 0000000..00ace9e --- /dev/null +++ b/tests/unit/cli/test_dataset_utils.py @@ -0,0 +1,174 @@ +"""Tests for CLI dataset_utils module.""" + +import pytest + +from airoa_lineage.cli.dataset_utils import ( + build_input_dataset, + build_output_dataset_common_facets, + parse_dataset_json, +) +from airoa_lineage.facets.dataset import ( + DestDirDatasetFacet, + OperationStatsDatasetFacet, + SrcDirDatasetFacet, + SrcFilesDatasetFacet, +) + + +class TestParseDatasetJson: + """Test parse_dataset_json function.""" + + def test_parse_valid_json(self): + """Test parsing valid JSON string.""" + json_string = '{"dataset_name": "test_dataset", "value": 123}' + result = parse_dataset_json(json_string, "input") + + assert result["dataset_name"] == "test_dataset" + assert result["value"] == 123 + + def test_parse_invalid_json(self): + """Test parsing invalid JSON raises ValueError.""" + json_string = '{"invalid": json}' + + with pytest.raises(ValueError) as exc_info: + parse_dataset_json(json_string, "input") + + assert "Invalid JSON in input dataset string" in str(exc_info.value) + + def test_parse_empty_json(self): + """Test parsing empty JSON object.""" + json_string = "{}" + result = parse_dataset_json(json_string, "output") + + assert result == {} + + +class TestBuildInputDataset: + """Test build_input_dataset function.""" + + def test_build_with_dataset_name_only(self): + """Test building dataset with only dataset_name.""" + data = {"dataset_name": "test_input"} + dataset = build_input_dataset(data, namespace="test_namespace") + + assert dataset.name == "test_input" + assert dataset.namespace == "test_namespace" + assert len(dataset.facets) == 0 + + def test_build_with_src_dir(self): + """Test building dataset with src_dir facet.""" + data = { + "dataset_name": "test_input", + "src_dir": {"path": "/data/source"}, + } + dataset = build_input_dataset(data, namespace="test_namespace") + + assert dataset.name == "test_input" + assert "srcDir" in dataset.facets + assert isinstance(dataset.facets["srcDir"], SrcDirDatasetFacet) + assert dataset.facets["srcDir"].path == "/data/source" + + def test_build_with_src_files(self): + """Test building dataset with src_files facet.""" + data = { + "dataset_name": "test_input", + "src_files": [ + {"name": "file1.txt", "hash": "abc123"}, + {"name": "file2.txt", "hash": "def456"}, + ], + } + dataset = build_input_dataset(data, namespace="test_namespace") + + assert dataset.name == "test_input" + assert "srcFiles" in dataset.facets + assert isinstance(dataset.facets["srcFiles"], SrcFilesDatasetFacet) + assert len(dataset.facets["srcFiles"].files) == 2 + assert dataset.facets["srcFiles"].files[0].name == "file1.txt" + assert dataset.facets["srcFiles"].files[0].hash == "abc123" + + def test_build_with_facet_prefix(self): + """Test building dataset with facet_prefix.""" + data = { + "dataset_name": "test_input", + "src_dir": {"path": "/data/source"}, + } + dataset = build_input_dataset( + data, namespace="test_namespace", facet_prefix="airoa_" + ) + + assert dataset.name == "test_input" + assert "airoa_srcDir" in dataset.facets + assert isinstance(dataset.facets["airoa_srcDir"], SrcDirDatasetFacet) + + def test_build_with_all_facets(self): + """Test building dataset with all facets.""" + data = { + "dataset_name": "test_input", + "src_dir": {"path": "/data/source"}, + "src_files": [ + {"name": "file1.txt", "hash": "abc123"}, + ], + } + dataset = build_input_dataset(data, namespace="test_namespace") + + assert dataset.name == "test_input" + assert "srcDir" in dataset.facets + assert "srcFiles" in dataset.facets + + +class TestBuildOutputDatasetCommonFacets: + """Test build_output_dataset_common_facets function.""" + + def test_build_with_empty_data(self): + """Test building facets with empty data.""" + data = {} + facets = build_output_dataset_common_facets(data) + + assert len(facets) == 0 + + def test_build_with_dest_dir(self): + """Test building facets with dest_dir.""" + data = { + "dest_dir": {"path": "/data/destination"}, + } + facets = build_output_dataset_common_facets(data) + + assert "destDir" in facets + assert isinstance(facets["destDir"], DestDirDatasetFacet) + assert facets["destDir"].path == "/data/destination" + + def test_build_with_operation_stats(self): + """Test building facets with operation_stats (new key).""" + data = { + "operation_stats": {"duration_seconds": 125.47}, + } + facets = build_output_dataset_common_facets(data) + + assert "operationStats" in facets + assert isinstance(facets["operationStats"], OperationStatsDatasetFacet) + assert facets["operationStats"].durationSeconds == 125.47 + + def test_build_with_facet_prefix(self): + """Test building facets with facet_prefix.""" + data = { + "dest_dir": {"path": "/data/destination"}, + "operation_stats": {"duration_seconds": 125.47}, + } + facets = build_output_dataset_common_facets(data, facet_prefix="airoa_") + + assert "airoa_destDir" in facets + assert "airoa_operationStats" in facets + assert isinstance(facets["airoa_destDir"], DestDirDatasetFacet) + assert isinstance(facets["airoa_operationStats"], OperationStatsDatasetFacet) + + def test_build_with_all_facets(self): + """Test building facets with all common facets.""" + data = { + "dest_dir": {"path": "/data/destination"}, + "operation_stats": {"duration_seconds": 125.47}, + } + facets = build_output_dataset_common_facets(data) + + assert len(facets) == 2 + assert "destDir" in facets + assert "operationStats" in facets diff --git a/tests/unit/cli/test_s3_upload.py b/tests/unit/cli/test_s3_upload.py index 25513c1..343cabc 100644 --- a/tests/unit/cli/test_s3_upload.py +++ b/tests/unit/cli/test_s3_upload.py @@ -2,10 +2,12 @@ import argparse - +from airoa_lineage.cli.common_args import ( + add_common_args, + add_common_facet_args, + add_repository_args, +) from airoa_lineage.cli.s3_upload import ( - _add_common_args, - _add_common_facet_args, create_parser, merge_config_with_args, ) @@ -15,9 +17,9 @@ class TestHelperFunctions: """Test helper functions for argument parsing.""" def test_add_common_args(self): - """Test that _add_common_args adds expected arguments.""" + """Test that add_common_args adds expected arguments.""" parser = argparse.ArgumentParser() - _add_common_args(parser) + add_common_args(parser) # Parse with arguments args = parser.parse_args( @@ -39,9 +41,10 @@ def test_add_common_args(self): assert args.facet_prefix == "test-prefix" def test_add_common_facet_args(self): - """Test that _add_common_facet_args adds expected arguments.""" + """Test that add_common_facet_args adds expected arguments.""" parser = argparse.ArgumentParser() - _add_common_facet_args(parser) + add_common_facet_args(parser) + add_repository_args(parser) # Parse with arguments args = parser.parse_args( diff --git a/tests/unit/cli/test_usb_copy.py b/tests/unit/cli/test_usb_copy.py index cd42aaf..1a64882 100644 --- a/tests/unit/cli/test_usb_copy.py +++ b/tests/unit/cli/test_usb_copy.py @@ -2,10 +2,13 @@ import argparse - +from airoa_lineage.cli.common_args import ( + add_common_args, + add_common_facet_args, + add_device_facet_args, + add_repository_args, +) from airoa_lineage.cli.usb_copy import ( - _add_common_args, - _add_common_facet_args, create_parser, merge_config_with_args, ) @@ -15,9 +18,9 @@ class TestHelperFunctions: """Test helper functions for argument parsing.""" def test_add_common_args(self): - """Test that _add_common_args adds expected arguments.""" + """Test that add_common_args adds expected arguments.""" parser = argparse.ArgumentParser() - _add_common_args(parser) + add_common_args(parser) # Parse with arguments args = parser.parse_args( @@ -39,9 +42,10 @@ def test_add_common_args(self): assert args.facet_prefix == "test-prefix" def test_add_common_facet_args(self): - """Test that _add_common_facet_args adds expected arguments.""" + """Test that add_common_facet_args adds expected arguments.""" parser = argparse.ArgumentParser() - _add_common_facet_args(parser) + add_common_facet_args(parser) + add_repository_args(parser) # Parse with arguments args = parser.parse_args( @@ -68,6 +72,21 @@ def test_add_common_facet_args(self): assert args.repository_tag == "v1.0.0" assert args.repository_branch == "main" + def test_add_device_facet_args(self): + """Test that add_device_facet_args adds expected arguments.""" + parser = argparse.ArgumentParser() + add_device_facet_args(parser) + + # Parse with arguments + args = parser.parse_args( + [ + "--hostname", + "test-host", + ] + ) + + assert args.hostname == "test-host" + class TestCreateParser: """Test argument parser creation.""" diff --git a/tests/unit/conversion/test_session.py b/tests/unit/conversion/test_session.py index 168a4da..c710994 100644 --- a/tests/unit/conversion/test_session.py +++ b/tests/unit/conversion/test_session.py @@ -422,7 +422,7 @@ def test_running_event_success(self, common_facet, aws_job_facet): aws_job_facet=aws_job_facet, ) session.start(input_datasets=[input_ds]) - session.running(message="Processing data") + session.running() # Verify emit was called twice (start + running) assert mock_client.emit.call_count == 2 @@ -448,7 +448,7 @@ def test_running_before_start_raises_error(self, common_facet, aws_job_facet): # running() without start() should raise error with pytest.raises(RuntimeError, match="not started"): - session.running(message="Should fail") + session.running() # Verify emit was never called mock_client.emit.assert_not_called() @@ -469,7 +469,7 @@ def test_running_after_complete_raises_error(self, common_facet, aws_job_facet): # running() after complete() should raise error with pytest.raises(RuntimeError, match="already completed"): - session.running(message="Should fail") + session.running() # Verify emit was called twice (start + complete, not running) assert mock_client.emit.call_count == 2 @@ -488,9 +488,9 @@ def test_multiple_running_events(self, common_facet, aws_job_facet): session.start(input_datasets=[input_ds]) # Call running() three times - session.running(message="Phase 1") - session.running(message="Phase 2") - session.running(message="Phase 3") + session.running() + session.running() + session.running() session.complete(output_datasets=[output_ds]) @@ -513,7 +513,7 @@ def test_running_event_includes_common_facet(self, common_facet, aws_job_facet): aws_job_facet=aws_job_facet, ) session.start(input_datasets=[input_ds]) - session.running(message="Processing") + session.running() # Check RUNNING event running_event = helper.get_emitted_event(mock_client, call_index=1) @@ -532,7 +532,7 @@ def test_running_event_includes_aws_job_facet(self, common_facet, aws_job_facet) aws_job_facet=aws_job_facet, ) session.start(input_datasets=[input_ds]) - session.running(message="Processing") + session.running() # Check RUNNING event running_event = helper.get_emitted_event(mock_client, call_index=1) @@ -552,7 +552,7 @@ def test_running_event_with_facet_prefix(self, common_facet, aws_job_facet): facet_prefix="airoa", ) session.start(input_datasets=[input_ds]) - session.running(message="Processing") + session.running() # Check RUNNING event running_event = helper.get_emitted_event(mock_client, call_index=1) @@ -609,7 +609,7 @@ def test_running_after_cancel_raises_error(self, common_facet, aws_job_facet): session.cancel() with pytest.raises(RuntimeError) as excinfo: - session.running(message="Should fail") + session.running() assert "was cancelled" in str(excinfo.value) assert "Cannot call running() after cancel()" in str(excinfo.value) diff --git a/tests/unit/facets/dataset/test_directory.py b/tests/unit/facets/dataset/test_directory.py new file mode 100644 index 0000000..0c659e0 --- /dev/null +++ b/tests/unit/facets/dataset/test_directory.py @@ -0,0 +1,41 @@ +"""Unit tests for SrcDirDatasetFacet and DestDirDatasetFacet.""" + +from airoa_lineage.facets.dataset import DestDirDatasetFacet, SrcDirDatasetFacet + + +def test_src_dir_facet_initialization(): + """Test SrcDirDatasetFacet initialization.""" + facet = SrcDirDatasetFacet(path="/media/usb0/robot_data/2025-12-08") + + assert facet.path == "/media/usb0/robot_data/2025-12-08" + + +def test_src_dir_facet_get_key(): + """Test SrcDirDatasetFacet.get_key() returns 'srcDir'.""" + assert SrcDirDatasetFacet.get_key() == "srcDir" + + +def test_src_dir_facet_with_absolute_path(): + """Test SrcDirDatasetFacet with absolute path.""" + facet = SrcDirDatasetFacet(path="/mnt/usb/data") + + assert facet.path == "/mnt/usb/data" + + +def test_dest_dir_facet_initialization(): + """Test DestDirDatasetFacet initialization.""" + facet = DestDirDatasetFacet(path="/data/robot_data/2025-12-08") + + assert facet.path == "/data/robot_data/2025-12-08" + + +def test_dest_dir_facet_get_key(): + """Test DestDirDatasetFacet.get_key() returns 'destDir'.""" + assert DestDirDatasetFacet.get_key() == "destDir" + + +def test_dest_dir_facet_with_absolute_path(): + """Test DestDirDatasetFacet with absolute path.""" + facet = DestDirDatasetFacet(path="/backup/robot_data") + + assert facet.path == "/backup/robot_data" diff --git a/tests/unit/facets/dataset/test_files.py b/tests/unit/facets/dataset/test_files.py new file mode 100644 index 0000000..1298e00 --- /dev/null +++ b/tests/unit/facets/dataset/test_files.py @@ -0,0 +1,73 @@ +"""Unit tests for SrcFilesDatasetFacet and SrcFileInfo.""" + +from airoa_lineage.facets.dataset import SrcFileInfo, SrcFilesDatasetFacet + + +def test_src_file_info_initialization(): + """Test SrcFileInfo initialization.""" + file_info = SrcFileInfo( + name="rosbag_001.bag", + hash="5d41402abc4b2a76b9719d911017c592", + ) + + assert file_info.name == "rosbag_001.bag" + assert file_info.hash == "5d41402abc4b2a76b9719d911017c592" + + +def test_src_files_facet_initialization(): + """Test SrcFilesDatasetFacet initialization.""" + files = [ + SrcFileInfo( + name="rosbag_001.bag", + hash="5d41402abc4b2a76b9719d911017c592", + ), + SrcFileInfo( + name="rosbag_002.bag", + hash="7d793037a0760186574b0282f2f435e7", + ), + ] + facet = SrcFilesDatasetFacet(files=files) + + assert len(facet.files) == 2 + assert facet.files[0].name == "rosbag_001.bag" + assert facet.files[0].hash == "5d41402abc4b2a76b9719d911017c592" + assert facet.files[1].name == "rosbag_002.bag" + assert facet.files[1].hash == "7d793037a0760186574b0282f2f435e7" + + +def test_src_files_facet_get_key(): + """Test SrcFilesDatasetFacet.get_key() returns 'srcFiles'.""" + assert SrcFilesDatasetFacet.get_key() == "srcFiles" + + +def test_src_files_facet_with_empty_list(): + """Test SrcFilesDatasetFacet with empty file list.""" + facet = SrcFilesDatasetFacet(files=[]) + + assert len(facet.files) == 0 + + +def test_src_files_facet_with_single_file(): + """Test SrcFilesDatasetFacet with single file.""" + files = [ + SrcFileInfo( + name="single.bag", + hash="abc123def456", + ), + ] + facet = SrcFilesDatasetFacet(files=files) + + assert len(facet.files) == 1 + assert facet.files[0].name == "single.bag" + assert facet.files[0].hash == "abc123def456" + + +def test_src_files_facet_with_multiple_files(): + """Test SrcFilesDatasetFacet with multiple files.""" + files = [SrcFileInfo(name=f"file_{i}.bag", hash=f"hash_{i}") for i in range(5)] + facet = SrcFilesDatasetFacet(files=files) + + assert len(facet.files) == 5 + for i in range(5): + assert facet.files[i].name == f"file_{i}.bag" + assert facet.files[i].hash == f"hash_{i}" diff --git a/tests/unit/facets/dataset/test_operation_stats.py b/tests/unit/facets/dataset/test_operation_stats.py new file mode 100644 index 0000000..4875b5f --- /dev/null +++ b/tests/unit/facets/dataset/test_operation_stats.py @@ -0,0 +1,36 @@ +"""Unit tests for OperationStatsDatasetFacet.""" + +from airoa_lineage.facets.dataset import OperationStatsDatasetFacet + + +def test_operation_stats_facet_initialization(): + """Test OperationStatsDatasetFacet initialization.""" + facet = OperationStatsDatasetFacet(durationSeconds=125.47) + + assert facet.durationSeconds == 125.47 + + +def test_operation_stats_facet_get_key(): + """Test OperationStatsDatasetFacet.get_key() returns 'operationStats'.""" + assert OperationStatsDatasetFacet.get_key() == "operationStats" + + +def test_operation_stats_facet_with_integer_duration(): + """Test OperationStatsDatasetFacet with integer duration.""" + facet = OperationStatsDatasetFacet(durationSeconds=60) + + assert facet.durationSeconds == 60 + + +def test_operation_stats_facet_with_fractional_duration(): + """Test OperationStatsDatasetFacet with fractional duration.""" + facet = OperationStatsDatasetFacet(durationSeconds=3.14159) + + assert facet.durationSeconds == 3.14159 + + +def test_operation_stats_facet_with_large_duration(): + """Test OperationStatsDatasetFacet with large duration value.""" + facet = OperationStatsDatasetFacet(durationSeconds=7200.5) # 2 hours 0.5 seconds + + assert facet.durationSeconds == 7200.5 diff --git a/tests/unit/facets/dataset/test_usb_device.py b/tests/unit/facets/dataset/test_usb_device.py new file mode 100644 index 0000000..61c7583 --- /dev/null +++ b/tests/unit/facets/dataset/test_usb_device.py @@ -0,0 +1,47 @@ +"""Unit tests for UsbDeviceDatasetFacet.""" + +from airoa_lineage.facets.dataset import UsbDeviceDatasetFacet + + +def test_usb_device_facet_initialization(): + """Test UsbDeviceDatasetFacet initialization.""" + facet = UsbDeviceDatasetFacet( + id="/dev/sdb1", + label="ROBOT_DATA_001", + fsType="ext4", + ) + + assert facet.id == "/dev/sdb1" + assert facet.label == "ROBOT_DATA_001" + assert facet.fsType == "ext4" + + +def test_usb_device_facet_get_key(): + """Test UsbDeviceDatasetFacet.get_key() returns 'usbDevice'.""" + assert UsbDeviceDatasetFacet.get_key() == "usbDevice" + + +def test_usb_device_facet_with_vfat(): + """Test UsbDeviceDatasetFacet with vfat filesystem.""" + facet = UsbDeviceDatasetFacet( + id="/dev/sdc1", + label="USB_BACKUP", + fsType="vfat", + ) + + assert facet.id == "/dev/sdc1" + assert facet.label == "USB_BACKUP" + assert facet.fsType == "vfat" + + +def test_usb_device_facet_with_exfat(): + """Test UsbDeviceDatasetFacet with exfat filesystem.""" + facet = UsbDeviceDatasetFacet( + id="/dev/sdd1", + label="EXFAT_DRIVE", + fsType="exfat", + ) + + assert facet.id == "/dev/sdd1" + assert facet.label == "EXFAT_DRIVE" + assert facet.fsType == "exfat" diff --git a/tests/unit/facets/run/__init__.py b/tests/unit/facets/run/__init__.py new file mode 100644 index 0000000..37d4c89 --- /dev/null +++ b/tests/unit/facets/run/__init__.py @@ -0,0 +1 @@ +"""Tests for Run Facets.""" diff --git a/tests/unit/facets/test_facet_keys.py b/tests/unit/facets/run/test_facet_keys.py similarity index 100% rename from tests/unit/facets/test_facet_keys.py rename to tests/unit/facets/run/test_facet_keys.py diff --git a/tests/unit/facets/run/test_job.py b/tests/unit/facets/run/test_job.py new file mode 100644 index 0000000..8ff6105 --- /dev/null +++ b/tests/unit/facets/run/test_job.py @@ -0,0 +1,27 @@ +"""Unit tests for JobRunFacet.""" + +from airoa_lineage.facets import JobRunFacet + + +def test_job_facet_initialization(): + """Test JobRunFacet initialization.""" + job_facet = JobRunFacet(id="test-job-123") + + assert job_facet.id == "test-job-123" + + +def test_job_facet_get_key(): + """Test JobRunFacet.get_key() returns 'job'.""" + assert JobRunFacet.get_key() == "job" + + +def test_job_facet_with_uuid(): + """Test JobRunFacet with UUID-formatted ID.""" + import uuid + + job_id = str(uuid.uuid4()) + job_facet = JobRunFacet(id=job_id) + + assert job_facet.id == job_id + # Verify it's a valid UUID format + assert uuid.UUID(job_facet.id) diff --git a/tests/unit/helpers/base_session_test.py b/tests/unit/helpers/base_session_test.py index d4926ce..dd97d87 100644 --- a/tests/unit/helpers/base_session_test.py +++ b/tests/unit/helpers/base_session_test.py @@ -15,6 +15,7 @@ class TestUSBCopySession(BaseSessionTest): PRODUCER_NAME = "airoa-usbcopy-system" """ +import inspect import os import uuid from typing import Type @@ -22,6 +23,7 @@ class TestUSBCopySession(BaseSessionTest): import pytest from airoa_lineage.core.base_session import BaseSession +from airoa_lineage.facets import JobRunFacet from tests.unit.helpers.test_helper import BaseSessionTestHelper @@ -46,12 +48,20 @@ def helper(self): # Initialization Tests # ======================================== - def test_initialization_with_defaults(self, common_facet): + def test_initialization_with_defaults(self, common_facet, device_facet): """Test initialization with default values.""" - session = self.SESSION_CLASS( - namespace="test_namespace", - common_facet=common_facet, - ) + # Check if SESSION_CLASS requires device_facet and job_facet + sig = inspect.signature(self.SESSION_CLASS.__init__) + requires_device_facet = "device_facet" in sig.parameters + requires_job_facet = "job_facet" in sig.parameters + + kwargs = {"namespace": "test_namespace", "common_facet": common_facet} + if requires_device_facet: + kwargs["device_facet"] = device_facet + if requires_job_facet: + kwargs["job_facet"] = JobRunFacet(id="test-job-123") + + session = self.SESSION_CLASS(**kwargs) assert session.namespace == "test_namespace" assert session.common_facet == common_facet @@ -71,13 +81,24 @@ def test_initialization_with_defaults(self, common_facet): assert session._started is False assert session._completed is False - def test_initialization_with_facet_prefix(self, common_facet): + def test_initialization_with_facet_prefix(self, common_facet, device_facet): """Test initialization with facet_prefix.""" - session = self.SESSION_CLASS( - namespace="test_namespace", - common_facet=common_facet, - facet_prefix="airoa", - ) + # Check if SESSION_CLASS requires device_facet and job_facet + sig = inspect.signature(self.SESSION_CLASS.__init__) + requires_device_facet = "device_facet" in sig.parameters + requires_job_facet = "job_facet" in sig.parameters + + kwargs = { + "namespace": "test_namespace", + "common_facet": common_facet, + "facet_prefix": "airoa", + } + if requires_device_facet: + kwargs["device_facet"] = device_facet + if requires_job_facet: + kwargs["job_facet"] = JobRunFacet(id="test-job-123") + + session = self.SESSION_CLASS(**kwargs) assert session.facet_prefix == "airoa" diff --git a/tests/unit/helpers/test_helper.py b/tests/unit/helpers/test_helper.py index 2abf03d..48aa787 100644 --- a/tests/unit/helpers/test_helper.py +++ b/tests/unit/helpers/test_helper.py @@ -6,7 +6,12 @@ from openlineage.client.run import RunState from airoa_lineage.core.base_session import BaseSession -from airoa_lineage.facets import AWSJobRunFacet, CommonRunFacet, DeviceRunFacet +from airoa_lineage.facets import ( + AWSJobRunFacet, + CommonRunFacet, + DeviceRunFacet, + JobRunFacet, +) class BaseSessionTestHelper: @@ -42,6 +47,7 @@ def create_session( self, namespace: str = "test_namespace", common_facet: Optional[CommonRunFacet] = None, + device_facet: Optional[DeviceRunFacet] = None, **kwargs: Any, ) -> Tuple[BaseSession, MagicMock]: """Create a session with mocked MarquezClient. @@ -49,6 +55,7 @@ def create_session( Args: namespace: Namespace for the session common_facet: CommonRunFacet instance + device_facet: DeviceRunFacet instance (auto-created if not provided) **kwargs: Additional arguments to pass to session constructor Returns: @@ -69,11 +76,31 @@ def create_session( self._mock_client = MagicMock() self._mock_client_class.return_value = self._mock_client - # Create session - merge common_facet into kwargs if provided + # Create session - merge common_facet and device_facet into kwargs if provided session_kwargs = dict(kwargs) if common_facet is not None: session_kwargs["common_facet"] = common_facet + # Auto-create device_facet if not provided and not already in kwargs + # Only add device_facet for sessions that require it (check __init__ signature) + import inspect + + sig = inspect.signature(self.session_class.__init__) + requires_device_facet = "device_facet" in sig.parameters + requires_job_facet = "job_facet" in sig.parameters + + if requires_device_facet and "device_facet" not in session_kwargs: + if device_facet is not None: + session_kwargs["device_facet"] = device_facet + else: + # Create default device_facet for sessions that require it + session_kwargs["device_facet"] = DeviceRunFacet(hostname="test-host") + + # Auto-create job_facet if not provided and not already in kwargs + if requires_job_facet and "job_facet" not in session_kwargs: + # Create default job_facet for sessions that require it + session_kwargs["job_facet"] = JobRunFacet(id="test-job-123") + # Create session (type ignore needed as subclasses accept different kwargs) session = self.session_class(namespace=namespace, **session_kwargs) # type: ignore[call-arg] @@ -206,6 +233,23 @@ def assert_aws_job_facet( assert facet.name == expected_facet.name assert facet.id == expected_facet.id + def assert_job_facet( + self, + event: Any, + expected_facet: "JobRunFacet", + facet_key: str = "job", + ) -> None: + """Assert that the event contains the expected JobRunFacet. + + Args: + event: The emitted event + expected_facet: Expected JobRunFacet instance + facet_key: Key in event.run.facets (default: "job") + """ + assert facet_key in event.run.facets + facet = event.run.facets[facet_key] + assert facet.id == expected_facet.id + def assert_nominal_time( self, event: Any, diff --git a/tests/unit/s3_upload/test_session.py b/tests/unit/s3_upload/test_session.py index d09615d..c54bc40 100644 --- a/tests/unit/s3_upload/test_session.py +++ b/tests/unit/s3_upload/test_session.py @@ -1,90 +1,474 @@ -"""Unit tests for S3UploadSession.""" +"""Unit tests for S3UploadSession with dataset support.""" +import os import uuid +import pytest +from openlineage.client.run import Dataset -from airoa_lineage.facets import CommonRunFacet from airoa_lineage.s3_upload import S3UploadSession -from tests.unit.helpers.base_session_test import BaseSessionTest +from tests.unit.helpers.test_helper import BaseSessionTestHelper -class TestS3UploadSessionInitialization(BaseSessionTest): +class TestS3UploadSessionInitialization: """Test S3UploadSession initialization.""" - SESSION_CLASS = S3UploadSession - DEFAULT_JOB_NAME = "s3-data-upload" - PRODUCER_NAME = "airoa-s3-system" + def test_initialization_with_defaults(self, common_facet, device_facet, job_facet): + """Test initialization with default values.""" + session = S3UploadSession( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) - def test_initialization_with_custom_values(self): + assert session.namespace == "test_namespace" + assert session.common_facet == common_facet + assert session.common_facet.robotId == "hsr001" + assert session.common_facet.location == "weblab" + assert session.common_facet.repositoryHash == "df110d5" + assert session.common_facet.repositoryUri == "https://github.com/user/repo.git" + assert session.common_facet.repositoryTag == "v1.0.0" + assert session.common_facet.repositoryBranch == "main" + assert session.device_facet == device_facet + assert session.device_facet.hostname == "copy-pc-001" + assert session.job_facet == job_facet + assert session.job_name == "s3-data-upload" + assert session.marquez_url == os.getenv("MARQUEZ_URL", "http://localhost:9000") + # run_id should be auto-generated UUID + assert session.run_id is not None + # Verify it's a valid UUID format + uuid.UUID(session.run_id) + # Session should not be started or completed + assert session._started is False + assert session._completed is False + # Input datasets should be empty list + assert session._input_datasets == [] + # Output datasets should be empty list + assert session._output_datasets == [] + + def test_default_job_name(self, common_facet, device_facet, job_facet): + """Test that default job_name is 's3-data-upload'.""" + session = S3UploadSession( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + + assert session.job_name == "s3-data-upload" + + def test_initialization_with_custom_values( + self, common_facet, device_facet, job_facet + ): """Test initialization with custom values.""" custom_run_id = str(uuid.uuid4()) - common_facet = CommonRunFacet( - robotId="hsr002", - location="lab_room_2", - repositoryHash="abc123", - repositoryUri="https://github.com/custom/repo.git", - repositoryTag="v2.0.0", - repositoryBranch="develop", - ) session = S3UploadSession( namespace="custom_namespace", common_facet=common_facet, - job_name="custom_job", + device_facet=device_facet, + job_facet=job_facet, + job_name="custom-s3-upload", marquez_url="http://example.com:9000", run_id=custom_run_id, ) assert session.namespace == "custom_namespace" - assert session.common_facet.robotId == "hsr002" - assert session.common_facet.location == "lab_room_2" - assert session.common_facet.repositoryHash == "abc123" - assert ( - session.common_facet.repositoryUri == "https://github.com/custom/repo.git" - ) - assert session.common_facet.repositoryTag == "v2.0.0" - assert session.common_facet.repositoryBranch == "develop" - assert session.job_name == "custom_job" + assert session.common_facet == common_facet + assert session.device_facet == device_facet + assert session.job_facet == job_facet + assert session.job_name == "custom-s3-upload" assert session.marquez_url == "http://example.com:9000" assert session.run_id == custom_run_id assert session._started is False assert session._completed is False -class TestS3UploadSessionStart(BaseSessionTest): +class TestS3UploadSessionStart: """Test S3UploadSession start() method.""" - SESSION_CLASS = S3UploadSession - DEFAULT_JOB_NAME = "s3-data-upload" - PRODUCER_NAME = "airoa-s3-system" + def test_start_returns_run_id(self, common_facet, device_facet, job_facet): + """Test that start() returns run_id.""" + with BaseSessionTestHelper(S3UploadSession) as helper: + input_ds = Dataset(namespace="test", name="input_dataset") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + run_id = session.start(input_datasets=[input_ds]) + + # Verify run_id is returned + assert run_id == session.run_id + # Verify session state changed + assert session._started is True + assert session._completed is False + # Verify input datasets were stored + assert session._input_datasets == [input_ds] + # Verify emit was called once + mock_client.emit.assert_called_once() + + def test_start_with_datasets(self, common_facet, device_facet, job_facet): + """Test that input datasets are included in START event (outputs are empty).""" + with BaseSessionTestHelper(S3UploadSession) as helper: + input_ds = Dataset(namespace="test", name="s3_input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + + # Verify input datasets are in the event + assert len(event.inputs) == 1 + assert event.inputs[0].namespace == "test" + assert event.inputs[0].name == "s3_input" + # Verify outputs are empty in START event + assert len(event.outputs) == 0 + + def test_start_twice_raises_error(self, common_facet, device_facet, job_facet): + """Test that calling start() twice raises RuntimeError.""" + with BaseSessionTestHelper(S3UploadSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + + # Second call should raise error + with pytest.raises(RuntimeError, match="already started"): + session.start(input_datasets=[input_ds]) + + # Verify emit was called only once (not twice) + assert mock_client.emit.call_count == 1 + + def test_start_with_nominal_time(self, common_facet, device_facet, job_facet): + """Test start() with nominal_start_time and nominal_end_time.""" + with BaseSessionTestHelper(S3UploadSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + run_id = session.start( + input_datasets=[input_ds], + nominal_start_time="2025-12-08T00:00:00+00:00", + nominal_end_time="2025-12-08T05:00:00+00:00", + ) + + # Verify run_id is returned + assert run_id == session.run_id + + # Verify emit was called with nominalTime facet + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + helper.assert_nominal_time( + event, + expected_start="2025-12-08T00:00:00+00:00", + expected_end="2025-12-08T05:00:00+00:00", + ) + + def test_start_includes_device_facet(self, common_facet, device_facet, job_facet): + """Test that START event includes device facet.""" + with BaseSessionTestHelper(S3UploadSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + + # Verify device facet is in the event + helper.assert_device_facet(event, device_facet) + + def test_start_includes_job_facet(self, common_facet, device_facet, job_facet): + """Test that START event includes job facet.""" + with BaseSessionTestHelper(S3UploadSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) -class TestS3UploadSessionComplete(BaseSessionTest): + # Verify job facet is in the event + helper.assert_job_facet(event, job_facet) + + def test_start_with_facet_prefix_includes_device( + self, common_facet, device_facet, job_facet + ): + """Test that START event includes device facet with custom prefix.""" + with BaseSessionTestHelper(S3UploadSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + facet_prefix="airoa", + ) + session.start(input_datasets=[input_ds]) + + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + + # Verify device facet is in the event with prefix + helper.assert_device_facet(event, device_facet, facet_key="airoa_device") + + +class TestS3UploadSessionComplete: """Test S3UploadSession complete() method.""" - SESSION_CLASS = S3UploadSession - DEFAULT_JOB_NAME = "s3-data-upload" - PRODUCER_NAME = "airoa-s3-system" + def test_complete_without_start_raises_error( + self, common_facet, device_facet, job_facet + ): + """Test that calling complete() without start() raises RuntimeError.""" + with BaseSessionTestHelper(S3UploadSession) as helper: + output_ds = Dataset(namespace="test", name="output") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + + # complete() without start() should raise error + with pytest.raises(RuntimeError, match="not started"): + session.complete(output_datasets=[output_ds]) + + # Verify emit was never called + mock_client.emit.assert_not_called() + + def test_complete_twice_raises_error(self, common_facet, device_facet, job_facet): + """Test that calling complete() twice raises RuntimeError.""" + with BaseSessionTestHelper(S3UploadSession) as helper: + input_ds = Dataset(namespace="test", name="input") + output_ds = Dataset(namespace="test", name="output") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + session.complete(output_datasets=[output_ds]) + + # Second call should raise error + with pytest.raises(RuntimeError, match="already completed"): + session.complete(output_datasets=[output_ds]) + + # Verify emit was called twice (start + complete, not 3 times) + assert mock_client.emit.call_count == 2 + + def test_complete_sets_completed_flag(self, common_facet, device_facet, job_facet): + """Test that complete() sets _completed flag.""" + with BaseSessionTestHelper(S3UploadSession) as helper: + input_ds = Dataset(namespace="test", name="input") + output_ds = Dataset(namespace="test", name="output") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + assert session._completed is False + + session.complete(output_datasets=[output_ds]) + assert session._completed is True + + def test_complete_with_datasets(self, common_facet, device_facet, job_facet): + """Test that output datasets are included in COMPLETE event.""" + with BaseSessionTestHelper(S3UploadSession) as helper: + input_ds = Dataset(namespace="test", name="s3_input") + output_ds = Dataset(namespace="test", name="s3_output") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + mock_client.emit.reset_mock() + + session.complete(output_datasets=[output_ds]) + + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + + # Verify output datasets are in the event + assert len(event.outputs) == 1 + assert event.outputs[0].namespace == "test" + assert event.outputs[0].name == "s3_output" + # Verify inputs are preserved + assert len(event.inputs) == 1 + assert event.inputs[0].name == "s3_input" + + +class TestS3UploadSessionDatasets: + """Test S3UploadSession dataset facets support.""" + + def test_start_with_dataset_facets(self, common_facet, device_facet, job_facet): + """Test start() with dataset facets (srcDir, srcFiles).""" + from airoa_lineage.facets.dataset import ( + SrcDirDatasetFacet, + SrcFileInfo, + SrcFilesDatasetFacet, + ) + + with BaseSessionTestHelper(S3UploadSession) as helper: + input_ds = Dataset( + namespace="test", + name="s3_input", + facets={ + "srcDir": SrcDirDatasetFacet(path="/data/robot_data"), + "srcFiles": SrcFilesDatasetFacet( + files=[ + SrcFileInfo(name="file1.txt", hash="abc123"), + SrcFileInfo(name="file2.txt", hash="def456"), + ] + ), + }, + ) + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + + # Verify input dataset facets + assert len(event.inputs) == 1 + input_dataset = event.inputs[0] + assert "srcDir" in input_dataset.facets + assert input_dataset.facets["srcDir"].path == "/data/robot_data" + assert "srcFiles" in input_dataset.facets + assert len(input_dataset.facets["srcFiles"].files) == 2 + + def test_complete_with_dataset_facets(self, common_facet, device_facet, job_facet): + """Test complete() with dataset facets (destDir, operationStats).""" + from airoa_lineage.facets.dataset import ( + DestDirDatasetFacet, + OperationStatsDatasetFacet, + ) + + with BaseSessionTestHelper(S3UploadSession) as helper: + input_ds = Dataset(namespace="test", name="s3_input") + output_ds = Dataset( + namespace="test", + name="s3_output", + facets={ + "destDir": DestDirDatasetFacet( + path="s3://my-bucket/robot_data/2025-12-08" + ), + "operationStats": OperationStatsDatasetFacet( + durationSeconds=325.89 + ), + }, + ) + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + # Start first + session.start(input_datasets=[input_ds]) + mock_client.emit.reset_mock() -class TestS3UploadSessionNominalTime(BaseSessionTest): - """Test S3UploadSession nominal time support.""" + # Complete with output dataset + session.complete(output_datasets=[output_ds]) - SESSION_CLASS = S3UploadSession - DEFAULT_JOB_NAME = "s3-data-upload" - PRODUCER_NAME = "airoa-s3-system" + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + # Verify output dataset facets + assert len(event.outputs) == 1 + output_dataset = event.outputs[0] + assert "destDir" in output_dataset.facets + assert ( + output_dataset.facets["destDir"].path + == "s3://my-bucket/robot_data/2025-12-08" + ) + assert "operationStats" in output_dataset.facets + assert output_dataset.facets["operationStats"].durationSeconds == 325.89 -class TestCommonFacet(BaseSessionTest): - """Test CommonRunFacet in S3UploadSession.""" + def test_start_with_nominal_time_and_datasets( + self, common_facet, device_facet, job_facet + ): + """Test start() with nominal time and datasets.""" + with BaseSessionTestHelper(S3UploadSession) as helper: + input_ds = Dataset(namespace="test", name="s3_input") - SESSION_CLASS = S3UploadSession - DEFAULT_JOB_NAME = "s3-data-upload" - PRODUCER_NAME = "airoa-s3-system" + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + run_id = session.start( + input_datasets=[input_ds], + nominal_start_time="2025-12-08T00:00:00+00:00", + nominal_end_time="2025-12-08T05:00:00+00:00", + ) + # Verify run_id is returned + assert run_id == session.run_id -class TestFacetPrefix(BaseSessionTest): - """Test facet_prefix functionality.""" + # Verify emit was called with nominalTime facet + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + helper.assert_nominal_time( + event, + expected_start="2025-12-08T00:00:00+00:00", + expected_end="2025-12-08T05:00:00+00:00", + ) - SESSION_CLASS = S3UploadSession - DEFAULT_JOB_NAME = "s3-data-upload" - PRODUCER_NAME = "airoa-s3-system" + # Verify datasets are also included + assert len(event.inputs) == 1 + assert event.inputs[0].name == "s3_input" diff --git a/tests/unit/teleop/test_session.py b/tests/unit/teleop/test_session.py index c461146..db25c7c 100644 --- a/tests/unit/teleop/test_session.py +++ b/tests/unit/teleop/test_session.py @@ -30,7 +30,7 @@ def test_initialization_with_defaults(self, common_facet, device_facet): assert session.common_facet.repositoryTag == "v1.0.0" assert session.common_facet.repositoryBranch == "main" assert session.device_facet == device_facet - assert session.device_facet.hostname == "operator-pc-001" + assert session.device_facet.hostname == "copy-pc-001" assert session.job_name == "robot-data-collection" assert session.marquez_url == os.getenv("MARQUEZ_URL", "http://localhost:9000") # run_id should be auto-generated UUID diff --git a/tests/unit/usb_copy/test_session.py b/tests/unit/usb_copy/test_session.py index 434917d..8d3ec6a 100644 --- a/tests/unit/usb_copy/test_session.py +++ b/tests/unit/usb_copy/test_session.py @@ -1,90 +1,569 @@ -"""Unit tests for USBCopySession.""" +"""Unit tests for USBCopySession with dataset support.""" +import os import uuid +import pytest +from openlineage.client.run import Dataset -from airoa_lineage.facets import CommonRunFacet from airoa_lineage.usb_copy import USBCopySession -from tests.unit.helpers.base_session_test import BaseSessionTest +from tests.unit.helpers.test_helper import BaseSessionTestHelper -class TestUSBCopySessionInitialization(BaseSessionTest): +class TestUSBCopySessionInitialization: """Test USBCopySession initialization.""" - SESSION_CLASS = USBCopySession - DEFAULT_JOB_NAME = "usb-data-copy" - PRODUCER_NAME = "airoa-usbcopy-system" + def test_initialization_with_defaults(self, common_facet, device_facet, job_facet): + """Test initialization with default values.""" + session = USBCopySession( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + + assert session.namespace == "test_namespace" + assert session.common_facet == common_facet + assert session.common_facet.robotId == "hsr001" + assert session.common_facet.location == "weblab" + assert session.common_facet.repositoryHash == "df110d5" + assert session.common_facet.repositoryUri == "https://github.com/user/repo.git" + assert session.common_facet.repositoryTag == "v1.0.0" + assert session.common_facet.repositoryBranch == "main" + assert session.device_facet == device_facet + assert session.device_facet.hostname == "copy-pc-001" + assert session.job_facet == job_facet + assert session.job_name == "usb-data-copy" + assert session.marquez_url == os.getenv("MARQUEZ_URL", "http://localhost:9000") + # run_id should be auto-generated UUID + assert session.run_id is not None + # Verify it's a valid UUID format + uuid.UUID(session.run_id) + # Session should not be started or completed + assert session._started is False + assert session._completed is False + # Input datasets should be empty list + assert session._input_datasets == [] + # Output datasets should be empty list + assert session._output_datasets == [] - def test_initialization_with_custom_values(self): + def test_default_job_name(self, common_facet, device_facet, job_facet): + """Test that default job_name is 'usb-data-copy'.""" + session = USBCopySession( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + + assert session.job_name == "usb-data-copy" + + def test_initialization_with_custom_values( + self, common_facet, device_facet, job_facet + ): """Test initialization with custom values.""" custom_run_id = str(uuid.uuid4()) - common_facet = CommonRunFacet( - robotId="hsr002", - location="lab_room_2", - repositoryHash="abc123", - repositoryUri="https://github.com/custom/repo.git", - repositoryTag="v2.0.0", - repositoryBranch="develop", - ) session = USBCopySession( namespace="custom_namespace", common_facet=common_facet, - job_name="custom_job", + device_facet=device_facet, + job_facet=job_facet, + job_name="custom-usb-copy", marquez_url="http://example.com:9000", run_id=custom_run_id, ) assert session.namespace == "custom_namespace" - assert session.common_facet.robotId == "hsr002" - assert session.common_facet.location == "lab_room_2" - assert session.common_facet.repositoryHash == "abc123" - assert ( - session.common_facet.repositoryUri == "https://github.com/custom/repo.git" - ) - assert session.common_facet.repositoryTag == "v2.0.0" - assert session.common_facet.repositoryBranch == "develop" - assert session.job_name == "custom_job" + assert session.common_facet == common_facet + assert session.device_facet == device_facet + assert session.job_facet == job_facet + assert session.job_name == "custom-usb-copy" assert session.marquez_url == "http://example.com:9000" assert session.run_id == custom_run_id assert session._started is False assert session._completed is False -class TestUSBCopySessionStart(BaseSessionTest): +class TestUSBCopySessionStart: """Test USBCopySession start() method.""" - SESSION_CLASS = USBCopySession - DEFAULT_JOB_NAME = "usb-data-copy" - PRODUCER_NAME = "airoa-usbcopy-system" + def test_start_returns_run_id(self, common_facet, device_facet, job_facet): + """Test that start() returns run_id.""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input_dataset") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + run_id = session.start(input_datasets=[input_ds]) + + # Verify run_id is returned + assert run_id == session.run_id + # Verify session state changed + assert session._started is True + assert session._completed is False + # Verify input datasets were stored + assert session._input_datasets == [input_ds] + # Verify emit was called once + mock_client.emit.assert_called_once() + + def test_start_with_datasets(self, common_facet, device_facet, job_facet): + """Test that input datasets are included in START event (outputs are empty).""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="usb_input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + + # Verify input datasets are in the event + assert len(event.inputs) == 1 + assert event.inputs[0].namespace == "test" + assert event.inputs[0].name == "usb_input" + # Verify outputs are empty in START event + assert len(event.outputs) == 0 + + def test_start_twice_raises_error(self, common_facet, device_facet, job_facet): + """Test that calling start() twice raises RuntimeError.""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + + # Second call should raise error + with pytest.raises(RuntimeError, match="already started"): + session.start(input_datasets=[input_ds]) + + # Verify emit was called only once (not twice) + assert mock_client.emit.call_count == 1 + + def test_start_with_nominal_time(self, common_facet, device_facet, job_facet): + """Test start() with nominal_start_time and nominal_end_time.""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input") + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + run_id = session.start( + input_datasets=[input_ds], + nominal_start_time="2025-12-08T00:00:00+00:00", + nominal_end_time="2025-12-08T05:00:00+00:00", + ) -class TestUSBCopySessionComplete(BaseSessionTest): + # Verify run_id is returned + assert run_id == session.run_id + + # Verify emit was called with nominalTime facet + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + helper.assert_nominal_time( + event, + expected_start="2025-12-08T00:00:00+00:00", + expected_end="2025-12-08T05:00:00+00:00", + ) + + def test_start_includes_device_facet(self, common_facet, device_facet, job_facet): + """Test that START event includes device facet.""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + + # Verify device facet is in the event + helper.assert_device_facet(event, device_facet) + + def test_start_includes_job_facet(self, common_facet, device_facet, job_facet): + """Test that START event includes job facet.""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + + # Verify job facet is in the event + helper.assert_job_facet(event, job_facet) + + def test_start_with_facet_prefix_includes_device( + self, common_facet, device_facet, job_facet + ): + """Test that START event includes device facet with custom prefix.""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + facet_prefix="airoa", + ) + session.start(input_datasets=[input_ds]) + + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + + # Verify device facet is in the event with prefix + helper.assert_device_facet(event, device_facet, facet_key="airoa_device") + + +class TestUSBCopySessionComplete: """Test USBCopySession complete() method.""" - SESSION_CLASS = USBCopySession - DEFAULT_JOB_NAME = "usb-data-copy" - PRODUCER_NAME = "airoa-usbcopy-system" + def test_complete_without_start_raises_error( + self, common_facet, device_facet, job_facet + ): + """Test that calling complete() without start() raises RuntimeError.""" + with BaseSessionTestHelper(USBCopySession) as helper: + output_ds = Dataset(namespace="test", name="output") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + + # complete() without start() should raise error + with pytest.raises(RuntimeError, match="not started"): + session.complete(output_datasets=[output_ds]) + + # Verify emit was never called + mock_client.emit.assert_not_called() + + def test_complete_twice_raises_error(self, common_facet, device_facet, job_facet): + """Test that calling complete() twice raises RuntimeError.""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input") + output_ds = Dataset(namespace="test", name="output") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + session.complete(output_datasets=[output_ds]) + + # Second complete() should raise error + with pytest.raises(RuntimeError, match="already completed"): + session.complete(output_datasets=[output_ds]) + + # Verify emit was called twice (start + complete, not 3 times) + assert mock_client.emit.call_count == 2 + + def test_complete_success(self, common_facet, device_facet, job_facet): + """Test successful complete() after start().""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input") + output_ds = Dataset(namespace="test", name="output") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + session.complete(output_datasets=[output_ds]) + # Verify session state + helper.assert_complete_success(session, mock_client) -class TestUSBCopySessionNominalTime(BaseSessionTest): - """Test USBCopySession nominal time support.""" + # Verify COMPLETE event includes datasets + complete_event = helper.get_emitted_event(mock_client, call_index=1) + assert len(complete_event.inputs) == 1 + assert complete_event.inputs[0].name == "input" + assert len(complete_event.outputs) == 1 + assert complete_event.outputs[0].name == "output" - SESSION_CLASS = USBCopySession - DEFAULT_JOB_NAME = "usb-data-copy" - PRODUCER_NAME = "airoa-usbcopy-system" + def test_complete_includes_device_facet( + self, common_facet, device_facet, job_facet + ): + """Test that COMPLETE event includes device facet.""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input") + output_ds = Dataset(namespace="test", name="output") + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + session.complete(output_datasets=[output_ds]) -class TestCommonFacet(BaseSessionTest): - """Test CommonRunFacet in USBCopySession.""" + # Verify COMPLETE event includes device facet + complete_event = helper.get_emitted_event(mock_client, call_index=1) + helper.assert_device_facet(complete_event, device_facet) - SESSION_CLASS = USBCopySession - DEFAULT_JOB_NAME = "usb-data-copy" - PRODUCER_NAME = "airoa-usbcopy-system" + def test_complete_includes_job_facet(self, common_facet, device_facet, job_facet): + """Test that COMPLETE event includes job facet.""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input") + output_ds = Dataset(namespace="test", name="output") + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + session.complete(output_datasets=[output_ds]) -class TestFacetPrefix(BaseSessionTest): + # Verify COMPLETE event includes job facet + complete_event = helper.get_emitted_event(mock_client, call_index=1) + helper.assert_job_facet(complete_event, job_facet) + + +class TestCommonFacet: + """Test common facet functionality.""" + + def test_start_includes_common_facet(self, common_facet, device_facet, job_facet): + """Test that START event includes common facet.""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + + # Verify common facet exists and has correct values + helper.assert_common_facet(event, common_facet) + + def test_complete_includes_common_facet( + self, common_facet, device_facet, job_facet + ): + """Test that COMPLETE event includes common facet.""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input") + output_ds = Dataset(namespace="test", name="output") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + session.complete(output_datasets=[output_ds]) + + # Verify emit was called twice + assert mock_client.emit.call_count == 2 + + # Check COMPLETE event (second call) + complete_event = helper.get_emitted_event(mock_client, call_index=1) + + # Verify common facet exists in COMPLETE event + helper.assert_common_facet(complete_event, common_facet) + + +class TestFacetPrefix: """Test facet_prefix functionality.""" - SESSION_CLASS = USBCopySession - DEFAULT_JOB_NAME = "usb-data-copy" - PRODUCER_NAME = "airoa-usbcopy-system" + def test_default_facet_prefix_empty(self, common_facet, device_facet, job_facet): + """Test that default facet_prefix is empty string.""" + session = USBCopySession( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + + # Verify facet_prefix defaults to empty string + assert session.facet_prefix == "" + + def test_start_without_prefix_uses_default_key( + self, common_facet, device_facet, job_facet + ): + """Test that START event uses 'common' key without prefix.""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + session.start(input_datasets=[input_ds]) + + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + + # Verify facet key is default (no prefix) + assert "common" in event.run.facets + helper.assert_common_facet(event, common_facet) + + def test_start_with_prefix_uses_prefixed_key( + self, common_facet, device_facet, job_facet + ): + """Test that START event uses prefixed key when facet_prefix is set.""" + with BaseSessionTestHelper(USBCopySession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + facet_prefix="airoa", + ) + session.start(input_datasets=[input_ds]) + + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + + # Verify facet key has prefix + assert "airoa_common" in event.run.facets + helper.assert_common_facet(event, common_facet, facet_key="airoa_common") + + # Verify unprefixed key does not exist + assert "common" not in event.run.facets + + +class TestUSBCopySessionCancel: + """Test USBCopySession cancel() method.""" + + def test_cancel_with_datasets(self, common_facet, device_facet, job_facet): + """Test that ABORT event includes input datasets.""" + helper = BaseSessionTestHelper(USBCopySession) + + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + + input_ds = Dataset(namespace="test_namespace", name="input_data") + + session.start(input_datasets=[input_ds]) + session.cancel() + + # Get the cancel event (second call) + cancel_event = helper.get_emitted_event(mock_client, call_index=1) + assert len(cancel_event.inputs) == 1 + assert cancel_event.inputs[0].name == "input_data" + # Output datasets are not set yet, so should be empty + assert len(cancel_event.outputs) == 0 + + def test_cancel_after_complete_raises_error( + self, common_facet, device_facet, job_facet + ): + """Test that calling cancel() after complete() raises RuntimeError.""" + helper = BaseSessionTestHelper(USBCopySession) + + with helper: + session, _ = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + + input_ds = Dataset(namespace="test_namespace", name="input_data") + output_ds = Dataset(namespace="test_namespace", name="output_data") + + session.start(input_datasets=[input_ds]) + session.complete(output_datasets=[output_ds]) + + with pytest.raises(RuntimeError) as excinfo: + session.cancel() + + assert "already completed" in str(excinfo.value) + + def test_cancel_includes_device_facet(self, common_facet, device_facet, job_facet): + """Test that ABORT event includes device facet.""" + helper = BaseSessionTestHelper(USBCopySession) + + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + + input_ds = Dataset(namespace="test_namespace", name="input_data") + session.start(input_datasets=[input_ds]) + session.cancel() + + cancel_event = helper.get_emitted_event(mock_client, call_index=1) + assert "common" in cancel_event.run.facets + assert "device" in cancel_event.run.facets + helper.assert_device_facet(cancel_event, device_facet) + + def test_cancel_includes_job_facet(self, common_facet, device_facet, job_facet): + """Test that ABORT event includes job facet.""" + helper = BaseSessionTestHelper(USBCopySession) + + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + job_facet=job_facet, + ) + + input_ds = Dataset(namespace="test_namespace", name="input_data") + session.start(input_datasets=[input_ds]) + session.cancel() + + cancel_event = helper.get_emitted_event(mock_client, call_index=1) + assert "job" in cancel_event.run.facets + helper.assert_job_facet(cancel_event, job_facet) diff --git a/tests/unit/usb_copy/test_session_datasets.py b/tests/unit/usb_copy/test_session_datasets.py new file mode 100644 index 0000000..30b39e3 --- /dev/null +++ b/tests/unit/usb_copy/test_session_datasets.py @@ -0,0 +1,180 @@ +"""Unit tests for USBCopySession dataset support.""" + +import pytest +from openlineage.client.run import Dataset + +from airoa_lineage.facets import CommonRunFacet +from airoa_lineage.facets.dataset import ( + DestDirDatasetFacet, + OperationStatsDatasetFacet, + SrcDirDatasetFacet, + SrcFileInfo, + SrcFilesDatasetFacet, + UsbDeviceDatasetFacet, +) +from airoa_lineage.usb_copy import USBCopySession +from tests.unit.helpers.test_helper import BaseSessionTestHelper + + +class TestUSBCopySessionDatasets: + """Test USBCopySession dataset support.""" + + @pytest.fixture + def helper(self): + """Create a test helper for USBCopySession.""" + helper = BaseSessionTestHelper(USBCopySession) + yield helper + helper.stop_patch() + + @pytest.fixture + def common_facet(self): + """Create a CommonRunFacet for testing.""" + return CommonRunFacet( + robotId="hsr001", + location="weblab", + repositoryHash="df110d5", + repositoryUri="https://github.com/user/repo.git", + repositoryTag="v1.0.0", + repositoryBranch="main", + ) + + def test_start_with_input_datasets(self, helper, common_facet): + """Test start() with input datasets.""" + # Create input dataset with USB device, source directory, and files facets + usb_facet = UsbDeviceDatasetFacet( + id="/dev/sdb1", + label="ROBOT_DATA_001", + fsType="ext4", + ) + src_dir_facet = SrcDirDatasetFacet(path="/media/usb0/robot_data/2025-12-08") + src_files_facet = SrcFilesDatasetFacet( + files=[ + SrcFileInfo(name="rosbag_001.bag", hash="abc123"), + SrcFileInfo(name="rosbag_002.bag", hash="def456"), + ] + ) + + input_ds = Dataset( + namespace="airoa_usb_copy", + name="usb_input_2025_12_08", + facets={ + "usbDevice": usb_facet, + "srcDir": src_dir_facet, + "srcFiles": src_files_facet, + }, + ) + + # Create session with helper + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + + # Start session with input datasets + run_id = session.start(input_datasets=[input_ds]) + + # Verify run_id returned + assert run_id == session.run_id + + # Verify input datasets were stored + assert len(session._input_datasets) == 1 + assert session._input_datasets[0] == input_ds + + # Verify emit was called once + mock_client.emit.assert_called_once() + + # Get emitted event + event = helper.get_emitted_event(mock_client, 0) + + # Verify event inputs + assert len(event.inputs) == 1 + assert event.inputs[0].namespace == "airoa_usb_copy" + assert event.inputs[0].name == "usb_input_2025_12_08" + + # Verify input dataset facets + input_facets = event.inputs[0].facets + assert "usbDevice" in input_facets + assert input_facets["usbDevice"].id == "/dev/sdb1" + assert input_facets["usbDevice"].label == "ROBOT_DATA_001" + assert input_facets["usbDevice"].fsType == "ext4" + + assert "srcDir" in input_facets + assert input_facets["srcDir"].path == "/media/usb0/robot_data/2025-12-08" + + assert "srcFiles" in input_facets + assert len(input_facets["srcFiles"].files) == 2 + assert input_facets["srcFiles"].files[0].name == "rosbag_001.bag" + assert input_facets["srcFiles"].files[0].hash == "abc123" + + def test_complete_with_output_datasets(self, helper, common_facet): + """Test complete() with output datasets.""" + # Create input dataset + input_ds = Dataset( + namespace="airoa_usb_copy", + name="usb_input_2025_12_08", + ) + + # Create output dataset with destination directory and operation stats facets + dest_dir_facet = DestDirDatasetFacet(path="/data/robot_data/2025-12-08") + operation_stats_facet = OperationStatsDatasetFacet(durationSeconds=125.47) + + output_ds = Dataset( + namespace="airoa_usb_copy", + name="usb_output_2025_12_08", + facets={ + "destDir": dest_dir_facet, + "operationStats": operation_stats_facet, + }, + ) + + # Create session with helper + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + + # Start session with input datasets + session.start(input_datasets=[input_ds]) + + # Complete session with output datasets + session.complete(output_datasets=[output_ds]) + + # Verify output datasets were stored + assert len(session._output_datasets) == 1 + assert session._output_datasets[0] == output_ds + + # Verify emit was called twice (start + complete) + assert mock_client.emit.call_count == 2 + + # Get COMPLETE event (second call) + event = helper.get_emitted_event(mock_client, 1) + + # Verify event outputs + assert len(event.outputs) == 1 + assert event.outputs[0].namespace == "airoa_usb_copy" + assert event.outputs[0].name == "usb_output_2025_12_08" + + # Verify output dataset facets + output_facets = event.outputs[0].facets + assert "destDir" in output_facets + assert output_facets["destDir"].path == "/data/robot_data/2025-12-08" + + assert "operationStats" in output_facets + assert output_facets["operationStats"].durationSeconds == 125.47 + + def test_complete_without_start_raises_error(self, helper, common_facet): + """Test complete() without start() raises RuntimeError.""" + output_ds = Dataset( + namespace="airoa_usb_copy", + name="usb_output_2025_12_08", + ) + + # Create session with helper + session, _ = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + + # Attempt to complete without starting should raise RuntimeError + with pytest.raises(RuntimeError, match="not started"): + session.complete(output_datasets=[output_ds])