diff --git a/README.md b/README.md index 5ed062e..841c4e0 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ 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` +- ☁️ **Wasabi Upload Management**: Track Wasabi cloud storage uploads with `WasabiUploadSession` - 🔌 **USB Data Copy Management**: Track USB data copy operations with `USBCopySession` - 📊 **OpenLineage Integration**: Full OpenLineage specification support (START/COMPLETE/RUNNING/FAIL) - 🏷️ **Custom Run Facets**: Robot metadata (robotId, location, repository info) via `CommonRunFacet` @@ -185,6 +186,14 @@ uv run python examples/teleop_session.py Demonstrates a complete teleoperation session with nominal time tracking. +### Wasabi Upload Session + +```bash +uv run python examples/wasabi_upload_session.py +``` + +Demonstrates a complete Wasabi cloud storage upload session with nominal time tracking. + ### Batch ETL Job ```bash diff --git a/docs/architecture.md b/docs/architecture.md index 10fa2ee..b8e6e97 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -93,6 +93,28 @@ ConversionSession.complete(output_datasets) # Emits COMPLETE event **API Documentation**: See docstrings in [session.py](../src/airoa_lineage/conversion/session.py) **Usage Example**: [examples/data_conversion.py](../examples/data_conversion.py) +### WasabiUploadSession + +**Location**: [src/airoa_lineage/wasabi_upload/session.py](../src/airoa_lineage/wasabi_upload/session.py) + +**Purpose**: Manages the lifecycle of Wasabi cloud storage upload operations with OpenLineage tracking. + +**Key Responsibilities**: +- Track Wasabi data upload sessions (START → COMPLETE) +- Automatic run_id generation and management +- Nominal time support for historical data processing +- State validation (prevents duplicate start/complete calls) + +**Lifecycle**: +``` +WasabiUploadSession.start() # Emits START event + └─> Data upload to Wasabi... +WasabiUploadSession.complete() # Emits COMPLETE event +``` + +**API Documentation**: See docstrings in [session.py](../src/airoa_lineage/wasabi_upload/session.py) +**Usage Example**: [examples/wasabi_upload_session.py](../examples/wasabi_upload_session.py) + ### USBCopySession **Location**: [src/airoa_lineage/usb_copy/session.py](../src/airoa_lineage/usb_copy/session.py) @@ -247,6 +269,7 @@ For detailed event schemas, see the [OpenLineage specification](https://openline - [MarquezClient](../src/airoa_lineage/marquez_client/client.py) - Event emission and REST API - [TeleopSession](../src/airoa_lineage/teleop/session.py) - Robot data collection lifecycle - [ConversionSession](../src/airoa_lineage/conversion/session.py) - Data conversion pipeline lifecycle +- [WasabiUploadSession](../src/airoa_lineage/wasabi_upload/session.py) - Wasabi cloud 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 @@ -257,6 +280,7 @@ For detailed event schemas, see the [OpenLineage specification](https://openline - [examples/teleop_session.py](../examples/teleop_session.py) - Complete teleoperation session with nominal time - [examples/data_conversion.py](../examples/data_conversion.py) - Data conversion pipeline with RUNNING events +- [examples/wasabi_upload_session.py](../examples/wasabi_upload_session.py) - Wasabi upload session with nominal time - [examples/usb_copy_session.py](../examples/usb_copy_session.py) - USB data copy session with nominal time - [examples/simple_batch_etl.py](../examples/simple_batch_etl.py) - Batch ETL job with input/output datasets diff --git a/docs/development.md b/docs/development.md index 5686fec..4b0ffe7 100644 --- a/docs/development.md +++ b/docs/development.md @@ -138,15 +138,42 @@ Tests are organized in the `tests/unit/` directory: ``` tests/ +├── conftest.py # Common fixtures (common_facet, device_facet, aws_job_facet) └── unit/ ├── marquez_client/ │ └── test_client.py ├── teleop/ │ └── test_session.py + ├── conversion/ + │ └── test_session.py + ├── wasabi_upload/ + │ └── test_session.py + ├── usb_copy/ + │ └── test_session.py + ├── core/ + │ └── test_base_session.py + ├── helpers/ # Test utilities + │ ├── test_helper.py # BaseSessionTestHelper class + │ ├── base_session_test.py # BaseSessionTest base class + │ └── test_test_helper.py # Tests for the helper itself └── utils/ └── test_timestamps.py ``` +**Test Helpers:** + +- `BaseSessionTestHelper`: Utility class for creating mocked sessions and asserting common patterns + - `create_session()`: Create session with mocked MarquezClient + - `get_emitted_event()`: Extract emitted event at specific call index + - `assert_start_success()`, `assert_complete_success()`: Verify session state + - `assert_common_facet()`, `assert_device_facet()`, `assert_aws_job_facet()`: Verify facets + - `assert_nominal_time()`, `assert_event_structure()`: Verify event structure + +- `BaseSessionTest`: Base class for session tests with common test methods + - Subclasses define `SESSION_CLASS`, `DEFAULT_JOB_NAME`, `PRODUCER_NAME` + - Provides 19 common test methods for initialization, start, complete, nominal time, facets + - Reduces test duplication and improves maintainability + ### Run All Tests ```bash @@ -217,40 +244,44 @@ addopts = [ def test_1(self): # Avoid ``` -2. **Follow AAA pattern (Arrange, Act, Assert):** +2. **Use BaseSessionTest for session tests:** ```python - def test_complete_success(self): - # Arrange - 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="operator-pc-001") - session = TeleopSession( - namespace="test", - common_facet=common_facet, - device_facet=device_facet - ) - session.start() - - # Act - session.complete() - - # Assert - assert session._completed is True + from tests.unit.helpers.base_session_test import BaseSessionTest + from airoa_lineage.teleop import TeleopSession + + class TestTeleopSessionInitialization(BaseSessionTest): + """Test TeleopSession initialization.""" + + SESSION_CLASS = TeleopSession + DEFAULT_JOB_NAME = "teleop-session" + PRODUCER_NAME = "airoa-teleop-system" + + # Common tests (19 methods) are inherited from BaseSessionTest + # Only add session-specific tests here + def test_custom_device_facet(self, common_facet, device_facet): + session = TeleopSession( + namespace="test", + common_facet=common_facet, + device_facet=device_facet + ) + assert session.device_facet == device_facet ``` -3. **Use mocking for external dependencies:** +3. **Use BaseSessionTestHelper for custom tests:** ```python - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_start_returns_run_id(self, mock_client_class): - mock_client = MagicMock() - mock_client_class.return_value = mock_client - # ... test code + from tests.unit.helpers.test_helper import BaseSessionTestHelper + + def test_complete_success(self, common_facet): + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + session.complete() + + # Use helper assertions + helper.assert_complete_success(session, mock_client) ``` 4. **Test both success and failure cases:** @@ -477,6 +508,7 @@ airoa-lineage/ │ ├── __init__.py │ └── timestamps.py # Timestamp utilities ├── tests/ +│ ├── conftest.py # Common fixtures │ └── unit/ # Unit tests (mirrors src/ structure) │ ├── marquez_client/ │ │ └── test_client.py @@ -484,6 +516,16 @@ airoa-lineage/ │ │ └── test_session.py │ ├── conversion/ │ │ └── test_session.py +│ ├── wasabi_upload/ +│ │ └── test_session.py +│ ├── usb_copy/ +│ │ └── test_session.py +│ ├── core/ +│ │ └── test_base_session.py +│ ├── helpers/ # Test utilities +│ │ ├── test_helper.py +│ │ ├── base_session_test.py +│ │ └── test_test_helper.py │ └── utils/ │ └── test_timestamps.py ├── .gitignore diff --git a/examples/wasabi_upload_session.py b/examples/wasabi_upload_session.py new file mode 100644 index 0000000..5a05c1f --- /dev/null +++ b/examples/wasabi_upload_session.py @@ -0,0 +1,187 @@ +""" +Wasabi Data Upload Session Example + +This example demonstrates a complete Wasabi data upload session: +1. Send START event when upload begins with nominal time period +2. Simulate data upload to Wasabi cloud storage +3. Send COMPLETE event when upload ends + +It shows how to: +- Use WasabiUploadSession class to manage upload session lifecycle +- Create CommonRunFacet with robot and repository information +- Automatically track run_id without manual management +- Specify robot_id and location to identify the robot and its location +- Specify repository information (hash, URI, tag, branch) for traceability +- Specify nominal time period (the time range of data being processed) +- Use facet_prefix to namespace custom facets (e.g., "airoa_common") +- Track a complete Wasabi data upload session with OpenLineage +""" + +import os +import time +from datetime import datetime, timedelta, timezone + +from airoa_lineage.facets import CommonRunFacet +from airoa_lineage.marquez_client import MarquezClient +from airoa_lineage.wasabi_upload import WasabiUploadSession + + +def main(): + """Run a complete Wasabi data upload session simulation.""" + + # Session configuration + marquez_url = os.getenv("MARQUEZ_URL", "http://localhost:9000") + namespace = "airoa_examples" + job_name = "wasabi-data-upload" + + # Robot and repository information + robot_id = "hsr001" # Robot identifier + location = "weblab" # Location identifier + repository_hash = "df110d5" # Git commit hash + repository_uri = "https://github.com/AIRoA/airoa-lineage.git" # Repository URL + repository_tag = "v1.0.0" # Git tag + repository_branch = "main" # Git branch + + # Calculate nominal time period (data time range) + # nominal_start: 1 day ago from now + # nominal_end: 5 hours after nominal_start + now = datetime.now(timezone.utc) + nominal_start = now - timedelta(days=1) + nominal_end = nominal_start + timedelta(hours=5) + + # Create common facet with robot and repository information + common_facet = CommonRunFacet( + robotId=robot_id, + location=location, + repositoryHash=repository_hash, + repositoryUri=repository_uri, + repositoryTag=repository_tag, + repositoryBranch=repository_branch, + ) + + # Create session instance + # This automatically generates a run_id and initializes the Marquez client + # Using facet_prefix="airoa" to namespace our custom facets + session = WasabiUploadSession( + namespace=namespace, + common_facet=common_facet, + job_name=job_name, + marquez_url=marquez_url, + facet_prefix="airoa", + ) + + print("=" * 60) + print("Wasabi Data Upload Session Example") + print("=" * 60) + print(f"Namespace: {namespace}") + print(f"Job Name: {job_name}") + print() + print("Robot & Repository Information:") + print(f" Robot ID: {robot_id}") + print(f" Location: {location}") + print(f" Repo Hash: {repository_hash}") + print(f" Repo URI: {repository_uri}") + print(f" Repo Tag: {repository_tag}") + print(f" Repo Branch: {repository_branch}") + print() + print(f"Run ID: {session.run_id}") + print(f"Marquez: {marquez_url}") + print() + print("Nominal Time Period (data time range):") + print(f" Start: {nominal_start.isoformat()}") + print(f" End: {nominal_end.isoformat()}") + print() + + # ========== START Event ========== + print("[1/4] Sending START event...") + + # Send START event with nominal time period + # This specifies the time range of the data being processed + session.start( + nominal_start_time=nominal_start.isoformat(), + nominal_end_time=nominal_end.isoformat(), + ) + + print("✓ START event sent successfully") + print() + + # ========== Simulate Data Upload ========== + print("[2/4] Simulating data upload to Wasabi...") + + # In a real scenario, this is where you would: + # - Prepare data files for upload + # - Connect to Wasabi S3-compatible API + # - Upload files to Wasabi bucket + # - Verify upload integrity + upload_duration = 3 # seconds + time.sleep(upload_duration) + + print(f"✓ Data upload completed ({upload_duration} seconds)") + print(" - In a real scenario, this would:") + print(" • Prepare data files for upload") + print(" • Connect to Wasabi S3-compatible API") + print(" • Upload files to Wasabi bucket") + print(" • Verify upload integrity") + print() + + # ========== COMPLETE Event ========== + print("[3/4] Sending COMPLETE event...") + + # Send COMPLETE event using the session + # This sends an OpenLineage COMPLETE event to Marquez + session.complete() + + print("✓ COMPLETE event sent successfully") + print() + + # ========== Query Lineage ========== + print("[4/4] Querying job information...") + time.sleep(1) # Give Marquez time to process + + try: + # Initialize client to query job information + client = MarquezClient(marquez_url, namespace=namespace) + + # Query job information + job = client.get_job(job_name) + print("✓ Job information retrieved:") + print(f" - Latest run state: {job.get('latestRun', {}).get('state', 'N/A')}") + print(f" - Job type: {job.get('type', 'N/A')}") + print() + + except Exception as e: + print(f"⚠ Failed to query job: {e}") + print(" (This is expected if Marquez is still processing the events)") + print() + + # ========== Summary ========== + print("=" * 60) + print("Session Summary") + print("=" * 60) + print("✓ Wasabi data upload session completed successfully") + print(f"✓ Robot ID: {robot_id}") + print(f"✓ Location: {location}") + print(f"✓ Repository: {repository_uri}") + print(f"✓ Commit: {repository_hash} ({repository_branch})") + print(f"✓ Tag: {repository_tag}") + print(f"✓ Run ID: {session.run_id}") + print() + print("Next steps:") + print("1. View lineage in Marquez Web UI:") + print( + f" {marquez_url.replace(':9000', ':3000')}/lineage/job/{namespace}/{job_name}" + ) + print() + print("2. Query job via API:") + print(f" curl {marquez_url}/api/v1/namespaces/{namespace}/jobs/{job_name}") + print() + print("3. Query this specific run:") + print( + f" curl {marquez_url}/api/v1/namespaces/{namespace}/jobs/{job_name}/runs/{session.run_id}" + ) + print() + print("=" * 60) + + +if __name__ == "__main__": + main() diff --git a/src/airoa_lineage/conversion/session.py b/src/airoa_lineage/conversion/session.py index 5a7e58e..159d81e 100644 --- a/src/airoa_lineage/conversion/session.py +++ b/src/airoa_lineage/conversion/session.py @@ -1,18 +1,15 @@ """Data conversion session management with OpenLineage tracking.""" -import os -import uuid -from typing import Any, Dict, List, Optional, Union, cast +from typing import Dict, List, Optional -from openlineage.client.facet import NominalTimeRunFacet -from openlineage.client.run import Dataset, Job, Run, RunEvent, RunState +from openlineage.client.facet import BaseFacet +from openlineage.client.run import Dataset +from airoa_lineage.core import BaseSession from airoa_lineage.facets import AWSJobRunFacet, CommonRunFacet -from airoa_lineage.marquez_client import MarquezClient -from airoa_lineage.utils.timestamps import get_event_timestamp -class ConversionSession: +class ConversionSession(BaseSession): """ Manages a data conversion session with OpenLineage tracking. @@ -113,39 +110,36 @@ def __init__( ... ) >>> # Facet keys will be "airoa_common" and "airoa_awsJob" """ - self.namespace = namespace + super().__init__(namespace, job_name, marquez_url, run_id, facet_prefix) self.common_facet = common_facet self.aws_job_facet = aws_job_facet - self.job_name = job_name - self.facet_prefix = facet_prefix - # Resolve marquez_url with proper type - self.marquez_url = cast( - str, marquez_url or os.getenv("MARQUEZ_URL") or "http://localhost:9000" - ) - self.run_id = run_id or str(uuid.uuid4()) - - # Initialize Marquez client - self.client = MarquezClient(self.marquez_url, namespace=namespace) - - # State management - self._started = False - self._completed = False self._input_datasets: List[Dataset] = [] + self._output_datasets: List[Dataset] = [] - def _get_facet_key(self, facet: Union[CommonRunFacet, AWSJobRunFacet]) -> str: - """ - Get facet dictionary key with optional prefix. + def _get_session_facets(self) -> Dict[str, BaseFacet]: + """Return session-specific facets.""" + return { + "common": self.common_facet, + "awsJob": self.aws_job_facet, + } - Args: - facet: Facet instance with get_key() method + def _get_producer(self) -> str: + """Return producer identifier.""" + return "airoa-conversion-system" - Returns: - Prefixed facet key (e.g., "airoa_awsJob") or base key (e.g., "awsJob") - """ - base_key = facet.get_key() - return f"{self.facet_prefix}_{base_key}" if self.facet_prefix else base_key + def _get_inputs(self) -> List[Dataset]: + """Return input datasets.""" + return self._input_datasets - def start( + def _get_outputs(self) -> List[Dataset]: + """Return output datasets.""" + return self._output_datasets + + def _supports_running_events(self) -> bool: + """Enable running() events for conversion sessions.""" + return True + + def start( # type: ignore[override] self, input_datasets: List[Dataset], nominal_start_time: Optional[str] = None, @@ -209,47 +203,15 @@ def start( ... nominal_end_time="2025-10-23T01:30:00+00:00" ... ) """ - if self._started: - raise RuntimeError( - f"Session {self.run_id} already started. " - "Cannot call start() multiple times." - ) - - # Store input datasets for later use in complete() + # Store input datasets self._input_datasets = input_datasets + # Call parent's start() method + return super().start(nominal_start_time, nominal_end_time) - # Build run facets - run_facets: Dict[str, Any] = {} - - # Add common and AWS job facets - run_facets[self._get_facet_key(self.common_facet)] = self.common_facet - run_facets[self._get_facet_key(self.aws_job_facet)] = self.aws_job_facet - - # Add nominal time facet if provided - if nominal_start_time or nominal_end_time: - run_facets["nominalTime"] = NominalTimeRunFacet( - nominalStartTime=nominal_start_time or get_event_timestamp(), - nominalEndTime=nominal_end_time, - ) - - # Create and send START event - # Note: outputs will be specified in COMPLETE event - event = RunEvent( - eventType=RunState.START, - eventTime=get_event_timestamp(), - run=Run(runId=self.run_id, facets=run_facets or {}), - job=Job(namespace=self.namespace, name=self.job_name), - inputs=input_datasets, - outputs=[], - producer="airoa-conversion-system", - ) - - self.client.emit(event) - self._started = True - - return self.run_id - - def complete(self, output_datasets: List[Dataset]) -> None: + def complete( # type: ignore[override] + self, + output_datasets: List[Dataset], + ) -> None: """ Send COMPLETE event to Marquez. @@ -291,38 +253,10 @@ def complete(self, output_datasets: List[Dataset]) -> None: >>> # Perform conversion... >>> session.complete(output_datasets=[output_ds]) """ - if not self._started: - raise RuntimeError( - f"Session {self.run_id} not started. " - "Must call start() before complete()." - ) - - if self._completed: - raise RuntimeError( - f"Session {self.run_id} already completed. " - "Cannot call complete() multiple times." - ) - - # Build run facets - run_facets: Dict[str, Any] = {} - - # Add common and AWS job facets - run_facets[self._get_facet_key(self.common_facet)] = self.common_facet - run_facets[self._get_facet_key(self.aws_job_facet)] = self.aws_job_facet - - # Create and send COMPLETE event - event = RunEvent( - eventType=RunState.COMPLETE, - eventTime=get_event_timestamp(), - run=Run(runId=self.run_id, facets=run_facets), - job=Job(namespace=self.namespace, name=self.job_name), - inputs=self._input_datasets, - outputs=output_datasets, - producer="airoa-conversion-system", - ) - - self.client.emit(event) - self._completed = True + # Store output datasets + self._output_datasets = output_datasets + # Call parent's complete() method + return super().complete() def running(self, message: Optional[str] = None) -> None: """ @@ -366,33 +300,6 @@ def running(self, message: Optional[str] = None) -> None: >>> session.running(message="Phase 2: Converting to LeRobot format") >>> session.complete(output_datasets=[output_ds]) """ - if not self._started: - raise RuntimeError( - f"Session {self.run_id} not started. " - "Must call start() before running()." - ) - - if self._completed: - raise RuntimeError( - f"Session {self.run_id} already completed. " - "Cannot call running() after complete()." - ) - - # Build run facets - run_facets: Dict[str, Any] = {} - - # Add common and AWS job facets - run_facets[self._get_facet_key(self.common_facet)] = self.common_facet - run_facets[self._get_facet_key(self.aws_job_facet)] = self.aws_job_facet - - # Create and send RUNNING event - # Note: inputs and outputs are not specified for RUNNING events - event = RunEvent( - eventType=RunState.RUNNING, - eventTime=get_event_timestamp(), - run=Run(runId=self.run_id, facets=run_facets), - job=Job(namespace=self.namespace, name=self.job_name), - producer="airoa-conversion-system", - ) - - self.client.emit(event) + # 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/core/__init__.py b/src/airoa_lineage/core/__init__.py new file mode 100644 index 0000000..0fcf806 --- /dev/null +++ b/src/airoa_lineage/core/__init__.py @@ -0,0 +1,5 @@ +"""Core base classes for session management.""" + +from airoa_lineage.core.base_session import BaseSession + +__all__ = ["BaseSession"] diff --git a/src/airoa_lineage/core/base_session.py b/src/airoa_lineage/core/base_session.py new file mode 100644 index 0000000..a43c0a5 --- /dev/null +++ b/src/airoa_lineage/core/base_session.py @@ -0,0 +1,452 @@ +"""Base session class for OpenLineage event management. + +This module provides an abstract base class that implements the Template Method +pattern for managing OpenLineage event lifecycles. It encapsulates common logic +for session initialization, event emission, and facet management. +""" + +import os +import uuid +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional, cast + +from openlineage.client.facet import BaseFacet, NominalTimeRunFacet +from openlineage.client.run import Dataset, Job, Run, RunEvent, RunState + +from airoa_lineage.marquez_client import MarquezClient +from airoa_lineage.utils.timestamps import get_event_timestamp + + +class BaseSession(ABC): + """ + Abstract base class for OpenLineage session management. + + This class implements the Template Method pattern, providing a common + framework for managing OpenLineage event lifecycles. Subclasses must + implement abstract methods to customize behavior. + + The session lifecycle follows this pattern: + 1. Initialization: Create session with namespace, job name, and configuration + 2. START event: Signal job start with optional nominal time + 3. RUNNING events (optional): Send progress updates during execution + 4. COMPLETE/FAIL event: Signal job completion + + Template Methods (public API): + - start(): Send START event + - complete(): Send COMPLETE event + - running(): Send RUNNING event (optional, subclass-dependent) + + Abstract Methods (must implement): + - _get_session_facets(): Return session-specific facets + - _get_producer(): Return producer identifier string + + Hook Methods (optional override): + - _get_inputs(): Return input datasets (default: []) + - _get_outputs(): Return output datasets (default: []) + - _supports_running_events(): Enable running() method (default: False) + + Examples: + >>> from abc import ABC + >>> from airoa_lineage.core.base_session import BaseSession + >>> from airoa_lineage.facets import CommonRunFacet + >>> + >>> class MySession(BaseSession): + ... def __init__(self, namespace, common_facet, **kwargs): + ... super().__init__(namespace, "my-job", **kwargs) + ... self.common_facet = common_facet + ... + ... def _get_session_facets(self): + ... return {"common": self.common_facet} + ... + ... def _get_producer(self): + ... return "my-system" + >>> + >>> common_facet = CommonRunFacet( + ... robotId="robot001", + ... location="lab", + ... repositoryHash="abc123", + ... repositoryUri="https://github.com/user/repo.git", + ... repositoryTag="v1.0.0", + ... repositoryBranch="main" + ... ) + >>> session = MySession("test_namespace", common_facet) + >>> run_id = session.start() + >>> # Do work... + >>> session.complete() + """ + + def __init__( + self, + namespace: str, + job_name: str, + marquez_url: Optional[str] = None, + run_id: Optional[str] = None, + facet_prefix: str = "", + ): + """ + Initialize a base session. + + Args: + namespace: OpenLineage namespace (required) + job_name: Job name (required) + 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}_{base_key}". + + Examples: + >>> # In subclass __init__ + >>> def __init__(self, namespace, common_facet, **kwargs): + ... super().__init__( + ... namespace=namespace, + ... job_name="my-job", + ... **kwargs + ... ) + ... self.common_facet = common_facet + """ + self.namespace = namespace + self.job_name = job_name + self.facet_prefix = facet_prefix + + # Resolve marquez_url with proper type + self.marquez_url = cast( + str, marquez_url or os.getenv("MARQUEZ_URL") or "http://localhost:9000" + ) + self.run_id = run_id or str(uuid.uuid4()) + + # Initialize Marquez client + self.client = MarquezClient(self.marquez_url, namespace=namespace) + + # State management + self._started = False + self._completed = False + + def _get_facet_key(self, facet: BaseFacet) -> str: + """ + Get facet dictionary key with optional prefix. + + Args: + facet: Facet instance with get_key() method + + Returns: + Prefixed facet key (e.g., "airoa_common") or base key (e.g., "common") + + Examples: + >>> # With prefix + >>> session = MySession(..., facet_prefix="airoa") + >>> key = session._get_facet_key(common_facet) + >>> print(key) # "airoa_common" + >>> + >>> # Without prefix + >>> session = MySession(..., facet_prefix="") + >>> key = session._get_facet_key(common_facet) + >>> print(key) # "common" + """ + base_key = facet.get_key() # type: ignore + return f"{self.facet_prefix}_{base_key}" if self.facet_prefix else base_key + + @abstractmethod + def _get_session_facets(self) -> Dict[str, BaseFacet]: + """ + Get session-specific facets (abstract method). + + Subclasses must implement this to return their custom facets. + + Returns: + Dictionary mapping facet keys to facet instances (without prefix). + The prefix will be applied automatically by _get_facet_key(). + + Examples: + >>> # Simple case: single facet + >>> def _get_session_facets(self): + ... return {"common": self.common_facet} + >>> + >>> # Complex case: multiple facets + >>> def _get_session_facets(self): + ... return { + ... "common": self.common_facet, + ... "device": self.device_facet + ... } + """ + pass + + @abstractmethod + def _get_producer(self) -> str: + """ + Get producer identifier (abstract method). + + Subclasses must implement this to return their producer string. + + Returns: + Producer identifier (e.g., "airoa-teleop-system") + + Examples: + >>> def _get_producer(self): + ... return "airoa-teleop-system" + """ + pass + + def _get_inputs(self) -> List[Dataset]: + """ + Get input datasets (hook method). + + Subclasses can override this to provide input datasets. + Default implementation returns empty list. + + Returns: + List of input Dataset objects + + Examples: + >>> # Override for dynamic inputs + >>> def _get_inputs(self): + ... return self._input_datasets + """ + return [] + + def _get_outputs(self) -> List[Dataset]: + """ + Get output datasets (hook method). + + Subclasses can override this to provide output datasets. + Default implementation returns empty list. + + Returns: + List of output Dataset objects + + Examples: + >>> # Override for dynamic outputs + >>> def _get_outputs(self): + ... return self._output_datasets + """ + return [] + + def _supports_running_events(self) -> bool: + """ + Check if running() events are supported (hook method). + + Subclasses can override this to enable running() method. + Default implementation returns False. + + Returns: + True if running() events are supported, False otherwise + + Examples: + >>> # Enable running events + >>> def _supports_running_events(self): + ... return True + """ + return False + + def _build_run_facets( + self, + nominal_start_time: Optional[str] = None, + nominal_end_time: Optional[str] = None, + ) -> Dict[str, Any]: + """ + Build run facets dictionary with prefix applied. + + This method collects facets from _get_session_facets() and applies + the facet_prefix to each key. Optionally adds nominal time facet. + + Args: + nominal_start_time: Start time of nominal period (ISO 8601) + nominal_end_time: End time of nominal period (ISO 8601) + + Returns: + Dictionary of run facets with prefixed keys + + Examples: + >>> # Internal use in start() method + >>> run_facets = self._build_run_facets( + ... nominal_start_time="2025-10-23T01:00:00+00:00", + ... nominal_end_time="2025-10-23T01:30:00+00:00" + ... ) + """ + run_facets: Dict[str, Any] = {} + + # Add session-specific facets with prefix + for base_key, facet in self._get_session_facets().items(): + prefixed_key = ( + f"{self.facet_prefix}_{base_key}" if self.facet_prefix else base_key + ) + run_facets[prefixed_key] = facet + + # Add nominal time facet if provided + if nominal_start_time or nominal_end_time: + run_facets["nominalTime"] = NominalTimeRunFacet( + nominalStartTime=nominal_start_time or get_event_timestamp(), + nominalEndTime=nominal_end_time, + ) + + return run_facets + + def _emit_event( + self, + event_type: RunState, + nominal_start_time: Optional[str] = None, + nominal_end_time: Optional[str] = None, + ) -> None: + """ + Emit OpenLineage event to Marquez. + + This method creates and sends an OpenLineage event with the specified + type, run facets, and input/output datasets. + + Args: + event_type: Event type (START, RUNNING, COMPLETE, FAIL) + nominal_start_time: Start time of nominal period (ISO 8601) + nominal_end_time: End time of nominal period (ISO 8601) + + Examples: + >>> # Internal use in template methods + >>> self._emit_event(RunState.START) + >>> self._emit_event(RunState.RUNNING) + >>> self._emit_event(RunState.COMPLETE) + """ + # Build run facets + run_facets = self._build_run_facets(nominal_start_time, nominal_end_time) + + # For RUNNING events, don't include inputs/outputs + if event_type == RunState.RUNNING: + inputs = [] + outputs = [] + else: + inputs = self._get_inputs() + outputs = self._get_outputs() + + # Create and send event + event = RunEvent( + eventType=event_type, + eventTime=get_event_timestamp(), + run=Run(runId=self.run_id, facets=run_facets), + job=Job(namespace=self.namespace, name=self.job_name), + inputs=inputs, + outputs=outputs, + producer=self._get_producer(), + ) + + self.client.emit(event) + + def start( + self, + nominal_start_time: Optional[str] = None, + nominal_end_time: Optional[str] = None, + ) -> str: + """ + Send START event to Marquez (template method). + + This method sends an OpenLineage START event to track the beginning + of the job. Optionally, you can specify the nominal time period + representing the data being processed. + + Args: + nominal_start_time: Start time of the data period being processed + (ISO 8601 format, e.g., "2025-10-23T01: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-10-23T01:30:00+00:00"). + Optional. + + Returns: + The run_id for this session + + Raises: + RuntimeError: If session was already started + + Examples: + >>> # Basic usage without nominal time + >>> session = MySession("test_namespace", common_facet) + >>> run_id = session.start() + >>> print(f"Session started: {run_id}") + >>> + >>> # With nominal time period + >>> run_id = session.start( + ... nominal_start_time="2025-10-23T01:00:00+00:00", + ... nominal_end_time="2025-10-23T01:30:00+00:00" + ... ) + """ + if self._started: + raise RuntimeError( + f"Session {self.run_id} already started. " + "Cannot call start() multiple times." + ) + + self._emit_event(RunState.START, nominal_start_time, nominal_end_time) + self._started = True + + return self.run_id + + def complete(self) -> None: + """ + Send COMPLETE event to Marquez (template method). + + This method sends an OpenLineage COMPLETE event to track the successful + completion of the job. + + Raises: + RuntimeError: If session was not started or already completed + + Examples: + >>> session = MySession("test_namespace", common_facet) + >>> session.start() + >>> # Do work... + >>> session.complete() + """ + if not self._started: + raise RuntimeError( + f"Session {self.run_id} not started. " + "Must call start() before complete()." + ) + + if self._completed: + raise RuntimeError( + f"Session {self.run_id} already completed. " + "Cannot call complete() multiple times." + ) + + self._emit_event(RunState.COMPLETE) + self._completed = True + + def running(self) -> None: + """ + Send RUNNING event to Marquez (template method). + + This method sends an OpenLineage RUNNING event to track job progress + during execution. Only supported if _supports_running_events() returns True. + + Raises: + RuntimeError: If running events are not supported + RuntimeError: If session was not started or already completed + + Examples: + >>> # In subclass that supports running events + >>> def _supports_running_events(self): + ... return True + >>> + >>> session = MySession("test_namespace", common_facet) + >>> session.start() + >>> session.running() # Send progress update + >>> # Continue work... + >>> session.running() # Send another update + >>> session.complete() + """ + if not self._supports_running_events(): + raise RuntimeError( + f"Session type {self.__class__.__name__} does not support " + "running() events. Override _supports_running_events() to enable." + ) + + if not self._started: + raise RuntimeError( + f"Session {self.run_id} not started. " + "Must call start() before running()." + ) + + if self._completed: + raise RuntimeError( + f"Session {self.run_id} already completed. " + "Cannot call running() after complete()." + ) + + self._emit_event(RunState.RUNNING) diff --git a/src/airoa_lineage/teleop/session.py b/src/airoa_lineage/teleop/session.py index 2ff29c5..b185988 100644 --- a/src/airoa_lineage/teleop/session.py +++ b/src/airoa_lineage/teleop/session.py @@ -1,18 +1,14 @@ """Teleoperation session management with OpenLineage tracking.""" -import os -import uuid -from typing import Any, Dict, Optional, Union, cast +from typing import Dict, Optional -from openlineage.client.facet import NominalTimeRunFacet -from openlineage.client.run import Job, Run, RunEvent, RunState +from openlineage.client.facet import BaseFacet +from airoa_lineage.core import BaseSession from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet -from airoa_lineage.marquez_client import MarquezClient -from airoa_lineage.utils.timestamps import get_event_timestamp -class TeleopSession: +class TeleopSession(BaseSession): """ Manages a teleoperation data collection session with OpenLineage tracking. @@ -111,185 +107,17 @@ def __init__( ... ) >>> # Facet keys will be "airoa_common" and "airoa_device" """ - self.namespace = namespace + super().__init__(namespace, job_name, marquez_url, run_id, facet_prefix) self.common_facet = common_facet self.device_facet = device_facet - self.job_name = job_name - self.facet_prefix = facet_prefix - # Resolve marquez_url with proper type - self.marquez_url = cast( - str, marquez_url or os.getenv("MARQUEZ_URL") or "http://localhost:9000" - ) - self.run_id = run_id or str(uuid.uuid4()) - # Initialize Marquez client - self.client = MarquezClient(self.marquez_url, namespace=namespace) + def _get_session_facets(self) -> Dict[str, BaseFacet]: + """Return session-specific facets.""" + return { + "common": self.common_facet, + "device": self.device_facet, + } - # State management - self._started = False - self._completed = False - - def _get_facet_key(self, facet: Union[CommonRunFacet, DeviceRunFacet]) -> str: - """ - Get facet dictionary key with optional prefix. - - Args: - facet: Facet instance with get_key() method - - Returns: - Prefixed facet key (e.g., "airoa_common") or base key (e.g., "common") - """ - base_key = facet.get_key() - return f"{self.facet_prefix}_{base_key}" if self.facet_prefix else base_key - - def start( - self, - 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 teleoperation session. Optionally, you can specify the nominal - time period representing the data being processed. - - Args: - nominal_start_time: Start time of the data period being processed - (ISO 8601 format, e.g., "2025-10-23T01: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-10-23T01:30:00+00:00"). - Optional. - - Returns: - The run_id for this session - - Raises: - RuntimeError: If session was already started - - Examples: - >>> # Basic usage without nominal time - >>> from airoa_lineage.teleop import CommonRunFacet, DeviceRunFacet - >>> 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="operator-pc-001") - >>> session = TeleopSession( - ... namespace="my_namespace", - ... common_facet=common_facet, - ... device_facet=device_facet - ... ) - >>> run_id = session.start() - >>> print(f"Session started: {run_id}") - - >>> # With nominal time period - >>> run_id = session.start( - ... nominal_start_time="2025-10-23T01:00:00+00:00", - ... nominal_end_time="2025-10-23T01:30:00+00:00" - ... ) - """ - if self._started: - raise RuntimeError( - f"Session {self.run_id} already started. " - "Cannot call start() multiple times." - ) - - # Build run facets - run_facets: Dict[str, Any] = {} - - # Add common and device facets - run_facets[self._get_facet_key(self.common_facet)] = self.common_facet - run_facets[self._get_facet_key(self.device_facet)] = self.device_facet - - # Add nominal time facet if provided - if nominal_start_time or nominal_end_time: - run_facets["nominalTime"] = NominalTimeRunFacet( - nominalStartTime=nominal_start_time or get_event_timestamp(), - nominalEndTime=nominal_end_time, - ) - - # Create and send START event - event = RunEvent( - eventType=RunState.START, - eventTime=get_event_timestamp(), - run=Run(runId=self.run_id, facets=run_facets or {}), - job=Job(namespace=self.namespace, name=self.job_name), - inputs=[], - outputs=[], - producer="airoa-teleoperation-system", - ) - - self.client.emit(event) - self._started = True - - return self.run_id - - def complete(self) -> None: - """ - Send COMPLETE event to Marquez. - - This method sends an OpenLineage COMPLETE event to track the successful - completion of the teleoperation session. - - Raises: - RuntimeError: If session was not started or already completed - - Examples: - >>> from airoa_lineage.teleop import CommonRunFacet, DeviceRunFacet - >>> 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="operator-pc-001") - >>> session = TeleopSession( - ... namespace="my_namespace", - ... common_facet=common_facet, - ... device_facet=device_facet - ... ) - >>> session.start() - >>> # Collect data... - >>> session.complete() - """ - if not self._started: - raise RuntimeError( - f"Session {self.run_id} not started. " - "Must call start() before complete()." - ) - - if self._completed: - raise RuntimeError( - f"Session {self.run_id} already completed. " - "Cannot call complete() multiple times." - ) - - # Build run facets - run_facets: Dict[str, Any] = {} - - # Add common and device facets - run_facets[self._get_facet_key(self.common_facet)] = self.common_facet - run_facets[self._get_facet_key(self.device_facet)] = self.device_facet - - # Create and send COMPLETE event - event = RunEvent( - eventType=RunState.COMPLETE, - eventTime=get_event_timestamp(), - run=Run(runId=self.run_id, facets=run_facets), - job=Job(namespace=self.namespace, name=self.job_name), - inputs=[], - outputs=[], - producer="airoa-teleoperation-system", - ) - - self.client.emit(event) - self._completed = True + def _get_producer(self) -> str: + """Return producer identifier.""" + return "airoa-teleop-system" diff --git a/src/airoa_lineage/usb_copy/session.py b/src/airoa_lineage/usb_copy/session.py index 6b91f63..cfb3119 100644 --- a/src/airoa_lineage/usb_copy/session.py +++ b/src/airoa_lineage/usb_copy/session.py @@ -1,18 +1,14 @@ """USB data copy session management with OpenLineage tracking.""" -import os -import uuid -from typing import Any, Dict, Optional, cast +from typing import Dict, Optional -from openlineage.client.facet import NominalTimeRunFacet -from openlineage.client.run import Job, Run, RunEvent, RunState +from openlineage.client.facet import BaseFacet +from airoa_lineage.core import BaseSession from airoa_lineage.facets import CommonRunFacet -from airoa_lineage.marquez_client import MarquezClient -from airoa_lineage.utils.timestamps import get_event_timestamp -class USBCopySession: +class USBCopySession(BaseSession): """ Manages a USB data copy session with OpenLineage tracking. @@ -104,180 +100,13 @@ def __init__( ... ) >>> # Facet key will be "airoa_common" """ - self.namespace = namespace + super().__init__(namespace, job_name, marquez_url, run_id, facet_prefix) self.common_facet = common_facet - self.job_name = job_name - self.facet_prefix = facet_prefix - # Resolve marquez_url with proper type - self.marquez_url = cast( - str, marquez_url or os.getenv("MARQUEZ_URL") or "http://localhost:9000" - ) - self.run_id = run_id or str(uuid.uuid4()) - # Initialize Marquez client - self.client = MarquezClient(self.marquez_url, namespace=namespace) + def _get_session_facets(self) -> Dict[str, BaseFacet]: + """Return session-specific facets.""" + return {"common": self.common_facet} - # State management - self._started = False - self._completed = False - - def _get_facet_key(self, facet: "CommonRunFacet") -> str: - """ - Get facet dictionary key with optional prefix. - - Args: - facet: Facet instance with get_key() method - - Returns: - Prefixed facet key (e.g., "airoa_common") or base key (e.g., "common") - """ - base_key = facet.get_key() - return f"{self.facet_prefix}_{base_key}" if self.facet_prefix else base_key - - def start( - self, - 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: - nominal_start_time: Start time of the data period being processed - (ISO 8601 format, e.g., "2025-10-23T01: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-10-23T01:30:00+00:00"). - Optional. - - Returns: - The run_id for this session - - Raises: - RuntimeError: If session was already started - - Examples: - >>> # Basic usage without nominal time - >>> from airoa_lineage.facets import CommonRunFacet - >>> from airoa_lineage.usb_copy import USBCopySession - >>> common_facet = CommonRunFacet( - ... robotId="hsr001", - ... location="weblab", - ... repositoryHash="df110d5", - ... repositoryUri="https://github.com/user/repo.git", - ... repositoryTag="v1.0.0", - ... repositoryBranch="main" - ... ) - >>> session = USBCopySession( - ... namespace="my_namespace", - ... common_facet=common_facet - ... ) - >>> run_id = session.start() - >>> print(f"Session started: {run_id}") - - >>> # With nominal time period - >>> run_id = session.start( - ... nominal_start_time="2025-10-23T01:00:00+00:00", - ... nominal_end_time="2025-10-23T01:30:00+00:00" - ... ) - """ - if self._started: - raise RuntimeError( - f"Session {self.run_id} already started. " - "Cannot call start() multiple times." - ) - - # Build run facets - run_facets: Dict[str, Any] = {} - - # Add common facet - run_facets[self._get_facet_key(self.common_facet)] = self.common_facet - - # Add nominal time facet if provided - if nominal_start_time or nominal_end_time: - run_facets["nominalTime"] = NominalTimeRunFacet( - nominalStartTime=nominal_start_time or get_event_timestamp(), - nominalEndTime=nominal_end_time, - ) - - # Create and send START event - event = RunEvent( - eventType=RunState.START, - eventTime=get_event_timestamp(), - run=Run(runId=self.run_id, facets=run_facets or {}), - job=Job(namespace=self.namespace, name=self.job_name), - inputs=[], - outputs=[], - producer="airoa-usbcopy-system", - ) - - self.client.emit(event) - self._started = True - - return self.run_id - - def complete(self) -> None: - """ - Send COMPLETE event to Marquez. - - This method sends an OpenLineage COMPLETE event to track the successful - completion of the USB data copy session. - - Raises: - RuntimeError: If session was not started or already completed - - Examples: - >>> from airoa_lineage.facets import CommonRunFacet - >>> from airoa_lineage.usb_copy import USBCopySession - >>> common_facet = CommonRunFacet( - ... robotId="hsr001", - ... location="weblab", - ... repositoryHash="df110d5", - ... repositoryUri="https://github.com/user/repo.git", - ... repositoryTag="v1.0.0", - ... repositoryBranch="main" - ... ) - >>> session = USBCopySession( - ... namespace="my_namespace", - ... common_facet=common_facet - ... ) - >>> session.start() - >>> # Copy data from USB... - >>> session.complete() - """ - if not self._started: - raise RuntimeError( - f"Session {self.run_id} not started. " - "Must call start() before complete()." - ) - - if self._completed: - raise RuntimeError( - f"Session {self.run_id} already completed. " - "Cannot call complete() multiple times." - ) - - # Build run facets - run_facets: Dict[str, Any] = {} - - # Add common facet - run_facets[self._get_facet_key(self.common_facet)] = self.common_facet - - # Create and send COMPLETE event - event = RunEvent( - eventType=RunState.COMPLETE, - eventTime=get_event_timestamp(), - run=Run(runId=self.run_id, facets=run_facets), - job=Job(namespace=self.namespace, name=self.job_name), - inputs=[], - outputs=[], - producer="airoa-usbcopy-system", - ) - - self.client.emit(event) - self._completed = True + def _get_producer(self) -> str: + """Return producer identifier.""" + return "airoa-usbcopy-system" diff --git a/src/airoa_lineage/wasabi_upload/__init__.py b/src/airoa_lineage/wasabi_upload/__init__.py new file mode 100644 index 0000000..0784a7e --- /dev/null +++ b/src/airoa_lineage/wasabi_upload/__init__.py @@ -0,0 +1,5 @@ +"""Wasabi data upload session management with OpenLineage tracking.""" + +from airoa_lineage.wasabi_upload.session import WasabiUploadSession + +__all__ = ["WasabiUploadSession"] diff --git a/src/airoa_lineage/wasabi_upload/session.py b/src/airoa_lineage/wasabi_upload/session.py new file mode 100644 index 0000000..49758d0 --- /dev/null +++ b/src/airoa_lineage/wasabi_upload/session.py @@ -0,0 +1,112 @@ +"""Wasabi data upload session management with OpenLineage tracking.""" + +from typing import Dict, Optional + +from openlineage.client.facet import BaseFacet + +from airoa_lineage.core import BaseSession +from airoa_lineage.facets import CommonRunFacet + + +class WasabiUploadSession(BaseSession): + """ + Manages a Wasabi data upload session with OpenLineage tracking. + + This class handles the lifecycle of a Wasabi cloud storage upload operation, + automatically managing the run_id and ensuring proper event sequencing. + + Examples: + >>> # Basic usage + >>> from airoa_lineage.facets import CommonRunFacet + >>> from airoa_lineage.wasabi_upload import WasabiUploadSession + >>> common_facet = CommonRunFacet( + ... robotId="hsr001", + ... location="weblab", + ... repositoryHash="df110d5", + ... repositoryUri="https://github.com/user/repo.git", + ... repositoryTag="v1.0.0", + ... repositoryBranch="main" + ... ) + >>> session = WasabiUploadSession( + ... namespace="my_namespace", + ... common_facet=common_facet + ... ) + >>> run_id = session.start() + >>> # Upload data to Wasabi... + >>> session.complete() + + >>> # Custom configuration + >>> session = WasabiUploadSession( + ... namespace="production", + ... common_facet=common_facet, + ... job_name="my-wasabi-upload-job", + ... run_id="existing-run-id-123" + ... ) + >>> session.start() + >>> session.complete() + """ + + def __init__( + self, + namespace: str, + common_facet: CommonRunFacet, + job_name: str = "wasabi-data-upload", + marquez_url: Optional[str] = None, + run_id: Optional[str] = None, + facet_prefix: str = "", + ): + """ + Initialize a Wasabi data upload session. + + Args: + namespace: OpenLineage namespace (required) + common_facet: Common metadata (required, includes robotId, location, repository info) + job_name: Job name (default: "wasabi-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". + + Examples: + >>> # Basic usage + >>> from airoa_lineage.facets import CommonRunFacet + >>> from airoa_lineage.wasabi_upload import WasabiUploadSession + >>> common_facet = CommonRunFacet( + ... robotId="hsr001", + ... location="weblab", + ... repositoryHash="df110d5", + ... repositoryUri="https://github.com/user/repo.git", + ... repositoryTag="v1.0.0", + ... repositoryBranch="main" + ... ) + >>> session = WasabiUploadSession( + ... namespace="my_namespace", + ... common_facet=common_facet + ... ) + + >>> # Custom configuration + >>> session = WasabiUploadSession( + ... namespace="my_namespace", + ... common_facet=common_facet, + ... job_name="my_job", + ... marquez_url="http://marquez.example.com:9000" + ... ) + + >>> # With facet prefix + >>> session = WasabiUploadSession( + ... namespace="my_namespace", + ... common_facet=common_facet, + ... facet_prefix="airoa" + ... ) + >>> # Facet key will be "airoa_common" + """ + super().__init__(namespace, job_name, marquez_url, run_id, facet_prefix) + self.common_facet = common_facet + + def _get_session_facets(self) -> Dict[str, BaseFacet]: + """Return session-specific facets.""" + return {"common": self.common_facet} + + def _get_producer(self) -> str: + """Return producer identifier.""" + return "airoa-wasabi-system" diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..13d2b7d --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,61 @@ +"""Shared pytest fixtures for all tests.""" + +import pytest + +from airoa_lineage.facets import AWSJobRunFacet, CommonRunFacet, DeviceRunFacet + + +@pytest.fixture +def common_facet(): + """Standard CommonRunFacet for testing. + + Returns a CommonRunFacet with: + - robotId: hsr001 + - location: weblab + - repositoryHash: df110d5 + - repositoryUri: https://github.com/user/repo.git + - repositoryTag: v1.0.0 + - repositoryBranch: main + """ + return CommonRunFacet( + robotId="hsr001", + location="weblab", + repositoryHash="df110d5", + repositoryUri="https://github.com/user/repo.git", + repositoryTag="v1.0.0", + repositoryBranch="main", + ) + + +@pytest.fixture +def alternative_common_facet(): + """Alternative CommonRunFacet for testing different values. + + Returns a CommonRunFacet with: + - robotId: hsr002 + - location: lab_room_2 + - Other fields differ from common_facet + """ + return CommonRunFacet( + robotId="hsr002", + location="lab_room_2", + repositoryHash="abc123", + repositoryUri="https://github.com/custom/repo.git", + repositoryTag="v2.0.0", + repositoryBranch="develop", + ) + + +@pytest.fixture +def device_facet(): + """Standard DeviceRunFacet for testing.""" + return DeviceRunFacet(hostname="operator-pc-001") + + +@pytest.fixture +def aws_job_facet(): + """Standard AWSJobRunFacet for testing.""" + return AWSJobRunFacet( + name="conversion-job-001", + id="lambda-12345", + ) diff --git a/tests/unit/conversion/test_session.py b/tests/unit/conversion/test_session.py index c272e99..0e7c4f3 100644 --- a/tests/unit/conversion/test_session.py +++ b/tests/unit/conversion/test_session.py @@ -2,35 +2,13 @@ import os import uuid -from unittest.mock import MagicMock, patch import pytest from openlineage.client.run import Dataset, RunState from airoa_lineage.conversion import ConversionSession -from airoa_lineage.facets import AWSJobRunFacet, CommonRunFacet - - -@pytest.fixture -def common_facet(): - """Fixture for common run facet.""" - return CommonRunFacet( - robotId="hsr001", - location="weblab", - repositoryHash="df110d5", - repositoryUri="https://github.com/airoa-org/rebake.git", - repositoryTag="v1.0.0", - repositoryBranch="main", - ) - - -@pytest.fixture -def aws_job_facet(): - """Fixture for AWS job run facet.""" - return AWSJobRunFacet( - name="conversion-job-001", - id="lambda-12345", - ) +from airoa_lineage.facets import AWSJobRunFacet +from tests.unit.helpers.test_helper import BaseSessionTestHelper class TestConversionSessionInitialization: @@ -53,10 +31,7 @@ def test_initialization_with_defaults(self, common_facet, aws_job_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/airoa-org/rebake.git" - ) + 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.aws_job_facet == aws_job_facet @@ -109,377 +84,276 @@ def test_initialization_with_custom_values(self, common_facet, aws_job_facet): class TestConversionSessionStart: """Test ConversionSession start() method.""" - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_start_returns_run_id(self, mock_client_class, common_facet, aws_job_facet): + def test_start_returns_run_id(self, common_facet, aws_job_facet): """Test that start() returns run_id.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input_dataset") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_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() - - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_start_with_datasets(self, mock_client_class, common_facet, aws_job_facet): + with BaseSessionTestHelper(ConversionSession) as helper: + input_ds = Dataset(namespace="test", name="input_dataset") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_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, aws_job_facet): """Test that input datasets are included in START event (outputs are empty).""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="raw_rosbag") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - session.start(input_datasets=[input_ds]) - - # Verify emit was called - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] - - # Verify input datasets are in the event - assert len(event.inputs) == 1 - assert event.inputs[0].namespace == "test" - assert event.inputs[0].name == "raw_rosbag" - # Verify outputs are empty in START event - assert len(event.outputs) == 0 - - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_start_twice_raises_error( - self, mock_client_class, common_facet, aws_job_facet - ): - """Test that calling start() twice raises RuntimeError.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - session.start(input_datasets=[input_ds]) - - # Second call should raise error - with pytest.raises(RuntimeError, match="already started"): + with BaseSessionTestHelper(ConversionSession) as helper: + input_ds = Dataset(namespace="test", name="raw_rosbag") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + ) session.start(input_datasets=[input_ds]) - # Verify emit was called only once (not twice) - assert mock_client.emit.call_count == 1 + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_start_with_nominal_time( - self, mock_client_class, common_facet, aws_job_facet - ): - """Test start() with nominal_start_time and nominal_end_time.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input") + # Verify input datasets are in the event + assert len(event.inputs) == 1 + assert event.inputs[0].namespace == "test" + assert event.inputs[0].name == "raw_rosbag" + # Verify outputs are empty in START event + assert len(event.outputs) == 0 - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - run_id = session.start( - input_datasets=[input_ds], - nominal_start_time="2025-10-23T01:00:00+00:00", - nominal_end_time="2025-10-23T01:30:00+00:00", - ) + def test_start_twice_raises_error(self, common_facet, aws_job_facet): + """Test that calling start() twice raises RuntimeError.""" + with BaseSessionTestHelper(ConversionSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + ) + session.start(input_datasets=[input_ds]) - # Verify run_id is returned - assert run_id == session.run_id + # Second call should raise error + with pytest.raises(RuntimeError, match="already started"): + session.start(input_datasets=[input_ds]) - # Verify emit was called with nominalTime facet - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] - assert "nominalTime" in event.run.facets - assert ( - event.run.facets["nominalTime"].nominalStartTime - == "2025-10-23T01:00:00+00:00" - ) - assert ( - event.run.facets["nominalTime"].nominalEndTime - == "2025-10-23T01:30:00+00:00" - ) + # Verify emit was called only once (not twice) + assert mock_client.emit.call_count == 1 - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_start_includes_aws_job_facet( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_start_with_nominal_time(self, common_facet, aws_job_facet): + """Test start() with nominal_start_time and nominal_end_time.""" + with BaseSessionTestHelper(ConversionSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + ) + run_id = session.start( + input_datasets=[input_ds], + nominal_start_time="2025-10-23T01:00:00+00:00", + nominal_end_time="2025-10-23T01:30: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-10-23T01:00:00+00:00", + expected_end="2025-10-23T01:30:00+00:00", + ) + + def test_start_includes_aws_job_facet(self, common_facet, aws_job_facet): """Test that START event includes AWS job facet.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - session.start(input_datasets=[input_ds]) + with BaseSessionTestHelper(ConversionSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + ) + session.start(input_datasets=[input_ds]) - # Verify emit was called - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) - # Verify AWS job facet is in the event - assert "awsJob" in event.run.facets - assert event.run.facets["awsJob"].name == "conversion-job-001" - assert event.run.facets["awsJob"].id == "lambda-12345" + # Verify AWS job facet is in the event + helper.assert_aws_job_facet(event, aws_job_facet) - @patch("airoa_lineage.conversion.session.MarquezClient") def test_start_with_facet_prefix_includes_aws_job( - self, mock_client_class, common_facet, aws_job_facet + self, common_facet, aws_job_facet ): """Test that START event includes AWS job facet with custom prefix.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - facet_prefix="airoa", - ) - session.start(input_datasets=[input_ds]) + with BaseSessionTestHelper(ConversionSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + facet_prefix="airoa", + ) + session.start(input_datasets=[input_ds]) - # Verify emit was called - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) - # Verify AWS job facet is in the event with prefix - assert "airoa_awsJob" in event.run.facets - assert event.run.facets["airoa_awsJob"].name == "conversion-job-001" - assert event.run.facets["airoa_awsJob"].id == "lambda-12345" + # Verify AWS job facet is in the event with prefix + helper.assert_aws_job_facet(event, aws_job_facet, facet_key="airoa_awsJob") class TestConversionSessionComplete: """Test ConversionSession complete() method.""" - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_complete_without_start_raises_error( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_complete_without_start_raises_error(self, common_facet, aws_job_facet): """Test that calling complete() without start() raises RuntimeError.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - output_ds = Dataset(namespace="test", name="output") + with BaseSessionTestHelper(ConversionSession) as helper: + output_ds = Dataset(namespace="test", name="output") - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + ) - # complete() without start() should raise error - with pytest.raises(RuntimeError, match="not started"): - session.complete(output_datasets=[output_ds]) + # 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() + # Verify emit was never called + mock_client.emit.assert_not_called() - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_complete_twice_raises_error( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_complete_twice_raises_error(self, common_facet, aws_job_facet): """Test that calling complete() twice raises RuntimeError.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client + with BaseSessionTestHelper(ConversionSession) 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, + aws_job_facet=aws_job_facet, + ) + session.start(input_datasets=[input_ds]) + session.complete(output_datasets=[output_ds]) - input_ds = Dataset(namespace="test", name="input") - output_ds = Dataset(namespace="test", name="output") + # Second complete() should raise error + with pytest.raises(RuntimeError, match="already completed"): + session.complete(output_datasets=[output_ds]) - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - session.start(input_datasets=[input_ds]) - session.complete(output_datasets=[output_ds]) + # Verify emit was called twice (start + complete, not 3 times) + assert mock_client.emit.call_count == 2 - # Second complete() should raise error - with pytest.raises(RuntimeError, match="already completed"): + def test_complete_success(self, common_facet, aws_job_facet): + """Test successful complete() after start().""" + with BaseSessionTestHelper(ConversionSession) 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, + aws_job_facet=aws_job_facet, + ) + session.start(input_datasets=[input_ds]) session.complete(output_datasets=[output_ds]) - # Verify emit was called twice (start + complete, not 3 times) - assert mock_client.emit.call_count == 2 + # Verify session state + helper.assert_complete_success(session, mock_client) - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_complete_success(self, mock_client_class, common_facet, aws_job_facet): - """Test successful complete() after start().""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client + # 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" - input_ds = Dataset(namespace="test", name="input") - output_ds = Dataset(namespace="test", name="output") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - session.start(input_datasets=[input_ds]) - session.complete(output_datasets=[output_ds]) - - # Verify session state - assert session._started is True - assert session._completed is True - # Verify emit was called twice (start + complete) - assert mock_client.emit.call_count == 2 - - # Verify COMPLETE event includes datasets - complete_event = mock_client.emit.call_args_list[1][0][0] - 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" - - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_complete_includes_aws_job_facet( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_complete_includes_aws_job_facet(self, common_facet, aws_job_facet): """Test that COMPLETE event includes AWS job facet.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input") - output_ds = Dataset(namespace="test", name="output") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - session.start(input_datasets=[input_ds]) - session.complete(output_datasets=[output_ds]) + with BaseSessionTestHelper(ConversionSession) 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, + aws_job_facet=aws_job_facet, + ) + session.start(input_datasets=[input_ds]) + session.complete(output_datasets=[output_ds]) - # Verify COMPLETE event includes AWS job facet - complete_event = mock_client.emit.call_args_list[1][0][0] - assert "awsJob" in complete_event.run.facets - assert complete_event.run.facets["awsJob"].name == "conversion-job-001" - assert complete_event.run.facets["awsJob"].id == "lambda-12345" + # Verify COMPLETE event includes AWS job facet + complete_event = helper.get_emitted_event(mock_client, call_index=1) + helper.assert_aws_job_facet(complete_event, aws_job_facet) class TestCommonFacet: """Test common facet functionality.""" - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_start_includes_common_facet( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_start_includes_common_facet(self, common_facet, aws_job_facet): """Test that START event includes common facet.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client + with BaseSessionTestHelper(ConversionSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + ) + session.start(input_datasets=[input_ds]) - input_ds = Dataset(namespace="test", name="input") + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - session.start(input_datasets=[input_ds]) - - # Verify emit was called - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] - - # Verify common facet exists and has correct values - assert "common" in event.run.facets - assert event.run.facets["common"].robotId == "hsr001" - assert event.run.facets["common"].location == "weblab" - assert event.run.facets["common"].repositoryHash == "df110d5" - assert ( - event.run.facets["common"].repositoryUri - == "https://github.com/airoa-org/rebake.git" - ) - assert event.run.facets["common"].repositoryTag == "v1.0.0" - assert event.run.facets["common"].repositoryBranch == "main" + # Verify common facet exists and has correct values + helper.assert_common_facet(event, common_facet) - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_complete_includes_common_facet( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_complete_includes_common_facet(self, common_facet, aws_job_facet): """Test that COMPLETE event includes common facet.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client + with BaseSessionTestHelper(ConversionSession) 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, + aws_job_facet=aws_job_facet, + ) + session.start(input_datasets=[input_ds]) + session.complete(output_datasets=[output_ds]) - input_ds = Dataset(namespace="test", name="input") - output_ds = Dataset(namespace="test", name="output") + # Verify emit was called twice + assert mock_client.emit.call_count == 2 - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_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 = mock_client.emit.call_args_list[1][0][0] - - # Verify common facet exists in COMPLETE event - assert "common" in complete_event.run.facets - assert complete_event.run.facets["common"].robotId == "hsr001" - assert complete_event.run.facets["common"].location == "weblab" - assert complete_event.run.facets["common"].repositoryHash == "df110d5" - assert ( - complete_event.run.facets["common"].repositoryUri - == "https://github.com/airoa-org/rebake.git" - ) - assert complete_event.run.facets["common"].repositoryTag == "v1.0.0" - assert complete_event.run.facets["common"].repositoryBranch == "main" + # 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.""" - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_default_facet_prefix_empty( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_default_facet_prefix_empty(self, common_facet, aws_job_facet): """Test that default facet_prefix is empty string.""" - mock_client = MagicMock() - mock_client_class.return_value = mock_client - session = ConversionSession( namespace="test_namespace", common_facet=common_facet, @@ -489,259 +363,205 @@ def test_default_facet_prefix_empty( # Verify facet_prefix defaults to empty string assert session.facet_prefix == "" - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_start_without_prefix_uses_default_key( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_start_without_prefix_uses_default_key(self, common_facet, aws_job_facet): """Test that START event uses 'common' key without prefix.""" - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - session.start(input_datasets=[input_ds]) + with BaseSessionTestHelper(ConversionSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + ) + session.start(input_datasets=[input_ds]) - # Verify emit was called - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] + # 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 - assert event.run.facets["common"].robotId == "hsr001" + # Verify facet key is default (no prefix) + assert "common" in event.run.facets + helper.assert_common_facet(event, common_facet) - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_start_with_prefix_uses_prefixed_key( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_start_with_prefix_uses_prefixed_key(self, common_facet, aws_job_facet): """Test that START event uses prefixed key when facet_prefix is set.""" - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - facet_prefix="airoa", - ) - session.start(input_datasets=[input_ds]) + with BaseSessionTestHelper(ConversionSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + facet_prefix="airoa", + ) + session.start(input_datasets=[input_ds]) - # Verify emit was called - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] + # 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 - assert event.run.facets["airoa_common"].robotId == "hsr001" + # 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 + # Verify unprefixed key does not exist + assert "common" not in event.run.facets class TestConversionSessionRunning: """Tests for ConversionSession.running() method.""" - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_running_event_success( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_running_event_success(self, common_facet, aws_job_facet): """Test successful RUNNING event emission.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - session.start(input_datasets=[input_ds]) - session.running(message="Processing data") + with BaseSessionTestHelper(ConversionSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + ) + session.start(input_datasets=[input_ds]) + session.running(message="Processing data") - # Verify emit was called twice (start + running) - assert mock_client.emit.call_count == 2 + # Verify emit was called twice (start + running) + assert mock_client.emit.call_count == 2 - # Check RUNNING event (second call) - running_event = mock_client.emit.call_args_list[1][0][0] + # Check RUNNING event (second call) + running_event = helper.get_emitted_event(mock_client, call_index=1) - # Verify event type is RUNNING - assert running_event.eventType == RunState.RUNNING + # Verify event type is RUNNING + assert running_event.eventType == RunState.RUNNING - # Verify inputs and outputs are empty lists (default when not specified) - assert running_event.inputs == [] - assert running_event.outputs == [] + # Verify inputs and outputs are empty lists (default when not specified) + assert running_event.inputs == [] + assert running_event.outputs == [] - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_running_before_start_raises_error( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_running_before_start_raises_error(self, common_facet, aws_job_facet): """Test that calling running() before start() raises RuntimeError.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client + with BaseSessionTestHelper(ConversionSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + ) - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - - # running() without start() should raise error - with pytest.raises(RuntimeError, match="not started"): - session.running(message="Should fail") + # running() without start() should raise error + with pytest.raises(RuntimeError, match="not started"): + session.running(message="Should fail") - # Verify emit was never called - mock_client.emit.assert_not_called() + # Verify emit was never called + mock_client.emit.assert_not_called() - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_running_after_complete_raises_error( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_running_after_complete_raises_error(self, common_facet, aws_job_facet): """Test that calling running() after complete() raises RuntimeError.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input") - output_ds = Dataset(namespace="test", name="output") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - session.start(input_datasets=[input_ds]) - session.complete(output_datasets=[output_ds]) + with BaseSessionTestHelper(ConversionSession) 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, + aws_job_facet=aws_job_facet, + ) + session.start(input_datasets=[input_ds]) + session.complete(output_datasets=[output_ds]) - # running() after complete() should raise error - with pytest.raises(RuntimeError, match="already completed"): - session.running(message="Should fail") + # running() after complete() should raise error + with pytest.raises(RuntimeError, match="already completed"): + session.running(message="Should fail") - # Verify emit was called twice (start + complete, not running) - assert mock_client.emit.call_count == 2 + # Verify emit was called twice (start + complete, not running) + assert mock_client.emit.call_count == 2 - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_multiple_running_events( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_multiple_running_events(self, common_facet, aws_job_facet): """Test that multiple running() calls are allowed.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input") - output_ds = Dataset(namespace="test", name="output") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - session.start(input_datasets=[input_ds]) + with BaseSessionTestHelper(ConversionSession) 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, + aws_job_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") + # Call running() three times + session.running(message="Phase 1") + session.running(message="Phase 2") + session.running(message="Phase 3") - session.complete(output_datasets=[output_ds]) + session.complete(output_datasets=[output_ds]) - # Verify emit was called 5 times (start + 3 running + complete) - assert mock_client.emit.call_count == 5 + # Verify emit was called 5 times (start + 3 running + complete) + assert mock_client.emit.call_count == 5 - # Verify all running events are RUNNING type - for i in range(1, 4): # Calls 1, 2, 3 are RUNNING events - event = mock_client.emit.call_args_list[i][0][0] - assert event.eventType == RunState.RUNNING + # Verify all running events are RUNNING type + for i in range(1, 4): # Calls 1, 2, 3 are RUNNING events + event = helper.get_emitted_event(mock_client, call_index=i) + assert event.eventType == RunState.RUNNING - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_running_event_includes_common_facet( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_running_event_includes_common_facet(self, common_facet, aws_job_facet): """Test that RUNNING event includes common facet.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - session.start(input_datasets=[input_ds]) - session.running(message="Processing") + with BaseSessionTestHelper(ConversionSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + ) + session.start(input_datasets=[input_ds]) + session.running(message="Processing") - # Check RUNNING event - running_event = mock_client.emit.call_args_list[1][0][0] + # Check RUNNING event + running_event = helper.get_emitted_event(mock_client, call_index=1) - # Verify common facet exists and has correct values - assert "common" in running_event.run.facets - assert running_event.run.facets["common"].robotId == "hsr001" - assert running_event.run.facets["common"].location == "weblab" + # Verify common facet exists and has correct values + helper.assert_common_facet(running_event, common_facet) - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_running_event_includes_aws_job_facet( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_running_event_includes_aws_job_facet(self, common_facet, aws_job_facet): """Test that RUNNING event includes AWS job facet.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - ) - session.start(input_datasets=[input_ds]) - session.running(message="Processing") + with BaseSessionTestHelper(ConversionSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + ) + session.start(input_datasets=[input_ds]) + session.running(message="Processing") - # Check RUNNING event - running_event = mock_client.emit.call_args_list[1][0][0] + # Check RUNNING event + running_event = helper.get_emitted_event(mock_client, call_index=1) - # Verify AWS job facet exists and has correct values - assert "awsJob" in running_event.run.facets - assert running_event.run.facets["awsJob"].name == "conversion-job-001" - assert running_event.run.facets["awsJob"].id == "lambda-12345" + # Verify AWS job facet exists and has correct values + helper.assert_aws_job_facet(running_event, aws_job_facet) - @patch("airoa_lineage.conversion.session.MarquezClient") - def test_running_event_with_facet_prefix( - self, mock_client_class, common_facet, aws_job_facet - ): + def test_running_event_with_facet_prefix(self, common_facet, aws_job_facet): """Test that RUNNING event respects facet_prefix.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - input_ds = Dataset(namespace="test", name="input") - - session = ConversionSession( - namespace="test_namespace", - common_facet=common_facet, - aws_job_facet=aws_job_facet, - facet_prefix="airoa", - ) - session.start(input_datasets=[input_ds]) - session.running(message="Processing") + with BaseSessionTestHelper(ConversionSession) as helper: + input_ds = Dataset(namespace="test", name="input") + + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + facet_prefix="airoa", + ) + session.start(input_datasets=[input_ds]) + session.running(message="Processing") - # Check RUNNING event - running_event = mock_client.emit.call_args_list[1][0][0] + # Check RUNNING event + running_event = helper.get_emitted_event(mock_client, call_index=1) - # Verify facet key has prefix - assert "airoa_common" in running_event.run.facets - assert running_event.run.facets["airoa_common"].robotId == "hsr001" + # Verify facet key has prefix + assert "airoa_common" in running_event.run.facets + helper.assert_common_facet( + running_event, common_facet, facet_key="airoa_common" + ) - # Verify unprefixed key does not exist - assert "common" not in running_event.run.facets + # Verify unprefixed key does not exist + assert "common" not in running_event.run.facets diff --git a/tests/unit/core/__init__.py b/tests/unit/core/__init__.py new file mode 100644 index 0000000..62f7145 --- /dev/null +++ b/tests/unit/core/__init__.py @@ -0,0 +1 @@ +"""Unit tests for core module.""" diff --git a/tests/unit/core/test_base_session.py b/tests/unit/core/test_base_session.py new file mode 100644 index 0000000..07fa3bd --- /dev/null +++ b/tests/unit/core/test_base_session.py @@ -0,0 +1,481 @@ +"""Unit tests for BaseSession.""" + +import os +import uuid +from typing import Dict, List +from unittest.mock import MagicMock, patch + +import pytest +from openlineage.client.facet import BaseFacet +from openlineage.client.run import Dataset + +from airoa_lineage.core.base_session import BaseSession +from airoa_lineage.facets import CommonRunFacet + + +class MinimalSession(BaseSession): + """Minimal concrete implementation of BaseSession for testing.""" + + def __init__( + self, + namespace: str, + common_facet: CommonRunFacet, + job_name: str = "test-job", + **kwargs, + ): + """Initialize minimal session.""" + super().__init__(namespace, job_name, **kwargs) + self.common_facet = common_facet + + def _get_session_facets(self) -> Dict[str, BaseFacet]: + """Return common facet.""" + return {"common": self.common_facet} + + def _get_producer(self) -> str: + """Return test producer.""" + return "test-producer" + + +class RunningSession(BaseSession): + """Session that supports running events.""" + + def __init__( + self, + namespace: str, + common_facet: CommonRunFacet, + job_name: str = "test-job", + **kwargs, + ): + """Initialize running-enabled session.""" + super().__init__(namespace, job_name, **kwargs) + self.common_facet = common_facet + + def _get_session_facets(self) -> Dict[str, BaseFacet]: + """Return common facet.""" + return {"common": self.common_facet} + + def _get_producer(self) -> str: + """Return test producer.""" + return "test-producer" + + def _supports_running_events(self) -> bool: + """Enable running events.""" + return True + + +class DatasetSession(BaseSession): + """Session with input/output datasets.""" + + def __init__( + self, + namespace: str, + common_facet: CommonRunFacet, + job_name: str = "test-job", + **kwargs, + ): + """Initialize dataset session.""" + super().__init__(namespace, job_name, **kwargs) + self.common_facet = common_facet + self._input_datasets: List[Dataset] = [] + self._output_datasets: List[Dataset] = [] + + def _get_session_facets(self) -> Dict[str, BaseFacet]: + """Return common facet.""" + return {"common": self.common_facet} + + def _get_producer(self) -> str: + """Return test producer.""" + return "test-producer" + + 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 + + +@pytest.fixture +def common_facet(): + """Fixture for common run facet.""" + return CommonRunFacet( + robotId="robot001", + location="lab", + repositoryHash="abc123", + repositoryUri="https://github.com/user/repo.git", + repositoryTag="v1.0.0", + repositoryBranch="main", + ) + + +class TestBaseSessionInitialization: + """Test BaseSession initialization.""" + + def test_initialization_with_defaults(self, common_facet): + """Test initialization with default values.""" + session = MinimalSession( + namespace="test_namespace", + common_facet=common_facet, + ) + + assert session.namespace == "test_namespace" + assert session.job_name == "test-job" + 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 + uuid.UUID(session.run_id) # Verify it's a valid UUID + assert session.facet_prefix == "" + # Session should not be started or completed + assert session._started is False + assert session._completed is False + + def test_initialization_with_custom_values(self, common_facet): + """Test initialization with custom values.""" + custom_run_id = str(uuid.uuid4()) + session = MinimalSession( + namespace="custom_namespace", + common_facet=common_facet, + job_name="custom-job", + marquez_url="http://example.com:9000", + run_id=custom_run_id, + ) + + assert session.namespace == "custom_namespace" + assert session.job_name == "custom-job" + 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 + + def test_initialization_with_facet_prefix(self, common_facet): + """Test initialization with facet_prefix.""" + session = MinimalSession( + namespace="test_namespace", + common_facet=common_facet, + facet_prefix="airoa", + ) + + assert session.facet_prefix == "airoa" + + +class TestBaseSessionStart: + """Test BaseSession start() method.""" + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_start_returns_run_id(self, mock_client_class, common_facet): + """Test that start() returns run_id.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = MinimalSession( + namespace="test_namespace", + common_facet=common_facet, + ) + run_id = session.start() + + # 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 emit was called once + mock_client.emit.assert_called_once() + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_start_twice_raises_error(self, mock_client_class, common_facet): + """Test that calling start() twice raises RuntimeError.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = MinimalSession( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + # Second call should raise error + with pytest.raises(RuntimeError, match="already started"): + session.start() + + # Verify emit was called only once + assert mock_client.emit.call_count == 1 + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_start_with_nominal_time(self, mock_client_class, common_facet): + """Test start() with nominal time.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = MinimalSession( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start( + nominal_start_time="2025-10-23T01:00:00+00:00", + nominal_end_time="2025-10-23T01:30:00+00:00", + ) + + # Get the event that was passed to emit() + event = mock_client.emit.call_args[0][0] + # Verify nominalTime facet exists + assert "nominalTime" in event.run.facets + assert ( + event.run.facets["nominalTime"].nominalStartTime + == "2025-10-23T01:00:00+00:00" + ) + assert ( + event.run.facets["nominalTime"].nominalEndTime + == "2025-10-23T01:30:00+00:00" + ) + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_start_event_has_correct_producer(self, mock_client_class, common_facet): + """Test that START event has correct producer.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = MinimalSession( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + # Get the event that was passed to emit() + event = mock_client.emit.call_args[0][0] + assert event.producer == "test-producer" + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_start_event_has_empty_inputs_outputs( + self, mock_client_class, common_facet + ): + """Test that START event has empty inputs and outputs by default.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = MinimalSession( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + # Get the event that was passed to emit() + event = mock_client.emit.call_args[0][0] + assert event.inputs == [] + assert event.outputs == [] + + +class TestBaseSessionComplete: + """Test BaseSession complete() method.""" + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_complete_without_start_raises_error(self, mock_client_class, common_facet): + """Test that calling complete() without start() raises RuntimeError.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = MinimalSession( + namespace="test_namespace", + common_facet=common_facet, + ) + + # complete() without start() should raise error + with pytest.raises(RuntimeError, match="not started"): + session.complete() + + # Verify emit was never called + mock_client.emit.assert_not_called() + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_complete_twice_raises_error(self, mock_client_class, common_facet): + """Test that calling complete() twice raises RuntimeError.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = MinimalSession( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + session.complete() + + # Second call should raise error + with pytest.raises(RuntimeError, match="already completed"): + session.complete() + + # Verify emit was called twice (start + complete once) + assert mock_client.emit.call_count == 2 + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_complete_success(self, mock_client_class, common_facet): + """Test successful complete.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = MinimalSession( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + session.complete() + + # Verify session state + assert session._started is True + assert session._completed is True + # Verify emit was called twice + assert mock_client.emit.call_count == 2 + + +class TestBaseSessionRunning: + """Test BaseSession running() method.""" + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_running_not_supported_by_default(self, mock_client_class, common_facet): + """Test that running() is not supported by default.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = MinimalSession( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + # running() should raise error for MinimalSession + with pytest.raises(RuntimeError, match="does not support running"): + session.running() + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_running_supported_when_enabled(self, mock_client_class, common_facet): + """Test that running() works when enabled.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = RunningSession( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + session.running() + + # Verify emit was called twice (start + running) + assert mock_client.emit.call_count == 2 + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_running_without_start_raises_error(self, mock_client_class, common_facet): + """Test that running() without start() raises error.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = RunningSession( + namespace="test_namespace", + common_facet=common_facet, + ) + + # running() without start() should raise error + with pytest.raises(RuntimeError, match="not started"): + session.running() + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_running_after_complete_raises_error(self, mock_client_class, common_facet): + """Test that running() after complete() raises error.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = RunningSession( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + session.complete() + + # running() after complete() should raise error + with pytest.raises(RuntimeError, match="already completed"): + session.running() + + +class TestFacetPrefix: + """Test facet_prefix functionality.""" + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_facet_prefix_applied(self, mock_client_class, common_facet): + """Test that facet_prefix is applied correctly.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = MinimalSession( + namespace="test_namespace", + common_facet=common_facet, + facet_prefix="airoa", + ) + session.start() + + # Get the event that was passed to emit() + event = mock_client.emit.call_args[0][0] + # Verify common facet has prefix + assert "airoa_common" in event.run.facets + assert event.run.facets["airoa_common"].robotId == "robot001" + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_no_prefix_uses_default_key(self, mock_client_class, common_facet): + """Test that no prefix uses default facet key.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = MinimalSession( + namespace="test_namespace", + common_facet=common_facet, + facet_prefix="", + ) + session.start() + + # Get the event that was passed to emit() + event = mock_client.emit.call_args[0][0] + # Verify common facet uses default key + assert "common" in event.run.facets + assert "airoa_common" not in event.run.facets + + +class TestDatasets: + """Test input/output dataset support.""" + + @patch("airoa_lineage.core.base_session.MarquezClient") + def test_datasets_in_events(self, mock_client_class, common_facet): + """Test that datasets are included in events.""" + # Setup mock + mock_client = MagicMock() + mock_client_class.return_value = mock_client + + session = DatasetSession( + namespace="test_namespace", + common_facet=common_facet, + ) + + # Add datasets + input_ds = Dataset(namespace="test_namespace", name="input_data") + output_ds = Dataset(namespace="test_namespace", name="output_data") + session._input_datasets = [input_ds] + session._output_datasets = [output_ds] + + session.start() + + # Get the event that was passed to emit() + event = mock_client.emit.call_args[0][0] + assert len(event.inputs) == 1 + assert event.inputs[0].name == "input_data" + assert len(event.outputs) == 1 + assert event.outputs[0].name == "output_data" diff --git a/tests/unit/helpers/__init__.py b/tests/unit/helpers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/helpers/base_session_test.py b/tests/unit/helpers/base_session_test.py new file mode 100644 index 0000000..05c5d7c --- /dev/null +++ b/tests/unit/helpers/base_session_test.py @@ -0,0 +1,361 @@ +"""Base test class for session tests. + +This module provides a base test class that contains common test methods +shared across all session test classes (USBCopy, WasabiUpload, Teleop, Conversion). + +Subclasses must define the following class attributes: +- SESSION_CLASS: The session class to test (e.g., USBCopySession) +- DEFAULT_JOB_NAME: The default job name for the session +- PRODUCER_NAME: The producer name for the session + +Example: + class TestUSBCopySession(BaseSessionTest): + SESSION_CLASS = USBCopySession + DEFAULT_JOB_NAME = "usb-data-copy" + PRODUCER_NAME = "airoa-usbcopy-system" +""" + +import os +import uuid +from typing import Type + +import pytest + +from airoa_lineage.core.base_session import BaseSession +from tests.unit.helpers.test_helper import BaseSessionTestHelper + + +class BaseSessionTest: + """Base test class for session tests. + + This class provides common test methods for all session types. + Subclasses must define SESSION_CLASS, DEFAULT_JOB_NAME, and PRODUCER_NAME. + """ + + # Subclasses must override these + SESSION_CLASS: Type[BaseSession] = None # type: ignore[assignment] + DEFAULT_JOB_NAME: str = "" + PRODUCER_NAME: str = "" + + @pytest.fixture + def helper(self): + """Create a test helper for the session class.""" + return BaseSessionTestHelper(self.SESSION_CLASS) + + # ======================================== + # Initialization Tests + # ======================================== + + def test_initialization_with_defaults(self, common_facet): + """Test initialization with default values.""" + session = self.SESSION_CLASS( + namespace="test_namespace", + common_facet=common_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.job_name == self.DEFAULT_JOB_NAME + 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 + + def test_initialization_with_facet_prefix(self, common_facet): + """Test initialization with facet_prefix.""" + session = self.SESSION_CLASS( + namespace="test_namespace", + common_facet=common_facet, + facet_prefix="airoa", + ) + + assert session.facet_prefix == "airoa" + + # ======================================== + # Start Tests + # ======================================== + + def test_start_returns_run_id(self, helper, common_facet): + """Test that start() returns run_id.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + run_id = session.start() + + # Verify run_id is returned + assert run_id == session.run_id + # Verify session state and emit + helper.assert_start_success(session, mock_client) + + def test_start_twice_raises_error(self, helper, common_facet): + """Test that calling start() twice raises RuntimeError.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + # Second call should raise error + with pytest.raises(RuntimeError, match="already started"): + session.start() + + # Verify emit was called only once (not twice) + assert mock_client.emit.call_count == 1 + + def test_start_event_has_correct_producer(self, helper, common_facet): + """Test that START event has correct producer.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + event = helper.get_emitted_event(mock_client) + assert event.producer == self.PRODUCER_NAME + + def test_start_event_has_empty_inputs_outputs(self, helper, common_facet): + """Test that START event has empty inputs and outputs.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + event = helper.get_emitted_event(mock_client) + assert event.inputs == [] + assert event.outputs == [] + + # ======================================== + # Complete Tests + # ======================================== + + def test_complete_without_start_raises_error(self, helper, common_facet): + """Test that calling complete() without start() raises RuntimeError.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + + # complete() without start() should raise error + with pytest.raises(RuntimeError, match="not started"): + session.complete() + + # Verify emit was never called + mock_client.emit.assert_not_called() + + def test_complete_twice_raises_error(self, helper, common_facet): + """Test that calling complete() twice raises RuntimeError.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + session.complete() + + # Second call should raise error + with pytest.raises(RuntimeError, match="already completed"): + session.complete() + + # Verify emit was called twice (start + complete once) + assert mock_client.emit.call_count == 2 + + def test_complete_sets_completed_flag(self, helper, common_facet): + """Test that complete() sets _completed flag.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + session.complete() + + # Verify session state + helper.assert_complete_success(session, mock_client) + + def test_complete_event_has_correct_producer(self, helper, common_facet): + """Test that COMPLETE event has correct producer.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + session.complete() + + # Get the COMPLETE event (second call) + event = helper.get_emitted_event(mock_client, call_index=1) + assert event.producer == self.PRODUCER_NAME + + def test_complete_event_has_empty_inputs_outputs(self, helper, common_facet): + """Test that COMPLETE event has empty inputs and outputs.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + session.complete() + + # Get the COMPLETE event (second call) + event = helper.get_emitted_event(mock_client, call_index=1) + assert event.inputs == [] + assert event.outputs == [] + + # ======================================== + # Nominal Time Tests + # ======================================== + + def test_start_with_both_nominal_times(self, helper, common_facet): + """Test start() with both nominal start and end times.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start( + nominal_start_time="2025-10-23T01:00:00+00:00", + nominal_end_time="2025-10-23T01:30:00+00:00", + ) + + event = helper.get_emitted_event(mock_client) + helper.assert_nominal_time( + event, + expected_start="2025-10-23T01:00:00+00:00", + expected_end="2025-10-23T01:30:00+00:00", + ) + + def test_start_with_only_nominal_end_time(self, helper, common_facet): + """Test start() with only nominal end time.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start(nominal_end_time="2025-10-23T01:30:00+00:00") + + event = helper.get_emitted_event(mock_client) + # Verify nominalTime facet exists + assert "nominalTime" in event.run.facets + # nominalStartTime should be auto-filled with current time + assert event.run.facets["nominalTime"].nominalStartTime is not None + helper.assert_nominal_time( + event, + expected_end="2025-10-23T01:30:00+00:00", + ) + + def test_start_without_nominal_time(self, helper, common_facet): + """Test start() without nominal time.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + event = helper.get_emitted_event(mock_client) + # Verify nominalTime facet does not exist + assert "nominalTime" not in event.run.facets + + # ======================================== + # Common Facet Tests + # ======================================== + + def test_start_includes_common_facet(self, helper, common_facet): + """Test that START event includes common facet.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + event = helper.get_emitted_event(mock_client) + # Verify common facet exists + assert "common" in event.run.facets + assert event.run.facets["common"].robotId == "hsr001" + assert event.run.facets["common"].location == "weblab" + + def test_complete_includes_common_facet(self, helper, common_facet): + """Test that COMPLETE event includes common facet.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + session.complete() + + # Get the COMPLETE event (second call) + event = helper.get_emitted_event(mock_client, call_index=1) + # Verify common facet exists + assert "common" in event.run.facets + assert event.run.facets["common"].robotId == "hsr001" + assert event.run.facets["common"].location == "weblab" + + # ======================================== + # Facet Prefix Tests + # ======================================== + + def test_facet_prefix_in_start_event(self, helper, common_facet): + """Test that facet_prefix is applied in START event.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + facet_prefix="airoa", + ) + session.start() + + event = helper.get_emitted_event(mock_client) + # Verify common facet has prefix + assert "airoa_common" in event.run.facets + assert event.run.facets["airoa_common"].robotId == "hsr001" + + def test_facet_prefix_in_complete_event(self, helper, common_facet): + """Test that facet_prefix is applied in COMPLETE event.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + facet_prefix="airoa", + ) + session.start() + session.complete() + + # Get the COMPLETE event (second call) + event = helper.get_emitted_event(mock_client, call_index=1) + # Verify common facet has prefix + assert "airoa_common" in event.run.facets + assert event.run.facets["airoa_common"].robotId == "hsr001" + + def test_no_prefix_uses_default_key(self, helper, common_facet): + """Test that no prefix uses default facet key.""" + with helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + facet_prefix="", + ) + session.start() + + event = helper.get_emitted_event(mock_client) + # Verify common facet uses default key + assert "common" in event.run.facets + assert "airoa_common" not in event.run.facets diff --git a/tests/unit/helpers/test_helper.py b/tests/unit/helpers/test_helper.py new file mode 100644 index 0000000..309ceef --- /dev/null +++ b/tests/unit/helpers/test_helper.py @@ -0,0 +1,295 @@ +"""Test helper utilities for BaseSession and its subclasses.""" + +from typing import Any, Optional, Tuple, Type +from unittest.mock import MagicMock, patch + +from openlineage.client.run import RunState + +from airoa_lineage.core.base_session import BaseSession +from airoa_lineage.facets import AWSJobRunFacet, CommonRunFacet, DeviceRunFacet + + +class BaseSessionTestHelper: + """Helper class for testing BaseSession and its subclasses. + + This class provides common utilities for: + - Creating mocked sessions + - Extracting emitted events + - Asserting common patterns (facets, state, event structure) + + Example: + helper = BaseSessionTestHelper(USBCopySession) + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + helper.assert_start_success(session, mock_client) + """ + + def __init__(self, session_class: Type[BaseSession]): + """Initialize helper for a specific session class. + + Args: + session_class: The session class to test (e.g., USBCopySession) + """ + self.session_class = session_class + self._patch: Optional[Any] = None + self._mock_client_class: Optional[MagicMock] = None + self._mock_client: Optional[MagicMock] = None + + def create_session( + self, + namespace: str = "test_namespace", + common_facet: Optional[CommonRunFacet] = None, + **kwargs: Any, + ) -> Tuple[BaseSession, MagicMock]: + """Create a session with mocked MarquezClient. + + Args: + namespace: Namespace for the session + common_facet: CommonRunFacet instance + **kwargs: Additional arguments to pass to session constructor + + Returns: + Tuple of (session instance, mocked MarquezClient) + + Example: + session, mock_client = helper.create_session( + namespace="test_ns", + common_facet=common_facet, + job_name="custom-job", + ) + """ + # Start patch + self._patch = patch("airoa_lineage.core.base_session.MarquezClient") + self._mock_client_class = self._patch.start() + + # Setup mock client + self._mock_client = MagicMock() + self._mock_client_class.return_value = self._mock_client + + # Create session - merge common_facet into kwargs if provided + session_kwargs = dict(kwargs) + if common_facet is not None: + session_kwargs["common_facet"] = common_facet + + # Create session (type ignore needed as subclasses accept different kwargs) + session = self.session_class(namespace=namespace, **session_kwargs) # type: ignore[call-arg] + + return session, self._mock_client + + def stop_patch(self) -> None: + """Stop the MarquezClient patch.""" + if self._patch: + self._patch.stop() + self._patch = None + + def get_emitted_event(self, mock_client: MagicMock, call_index: int = 0) -> Any: + """Get the event emitted at a specific call index. + + Args: + mock_client: The mocked MarquezClient + call_index: Index of the emit call (0 = first call, 1 = second, etc.) + + Returns: + The emitted event object + + Example: + # Get the START event (first call) + start_event = helper.get_emitted_event(mock_client, 0) + + # Get the COMPLETE event (second call) + complete_event = helper.get_emitted_event(mock_client, 1) + """ + return mock_client.emit.call_args_list[call_index][0][0] + + def assert_start_success( + self, + session: BaseSession, + mock_client: MagicMock, + ) -> None: + """Assert that start() was successful. + + Verifies: + - Session state is _started=True, _completed=False + - MarquezClient.emit was called once + + Args: + session: The session instance + mock_client: The mocked MarquezClient + """ + assert session._started is True + assert session._completed is False + mock_client.emit.assert_called_once() + + def assert_complete_success( + self, + session: BaseSession, + mock_client: MagicMock, + ) -> None: + """Assert that complete() was successful. + + Verifies: + - Session state is _started=True, _completed=True + - MarquezClient.emit was called twice (start + complete) + + Args: + session: The session instance + mock_client: The mocked MarquezClient + """ + assert session._started is True + assert session._completed is True + assert mock_client.emit.call_count == 2 + + def assert_common_facet( + self, + event: Any, + expected_facet: CommonRunFacet, + facet_key: str = "common", + ) -> None: + """Assert that the event contains the expected CommonRunFacet. + + Args: + event: The emitted event + expected_facet: Expected CommonRunFacet instance + facet_key: Key in event.run.facets (default: "common") + + Example: + event = helper.get_emitted_event(mock_client) + helper.assert_common_facet(event, common_facet) + + # With custom prefix + helper.assert_common_facet(event, common_facet, facet_key="airoa_common") + """ + assert facet_key in event.run.facets + facet = event.run.facets[facet_key] + assert facet.robotId == expected_facet.robotId + assert facet.location == expected_facet.location + assert facet.repositoryHash == expected_facet.repositoryHash + assert facet.repositoryUri == expected_facet.repositoryUri + assert facet.repositoryTag == expected_facet.repositoryTag + assert facet.repositoryBranch == expected_facet.repositoryBranch + + def assert_device_facet( + self, + event: Any, + expected_facet: DeviceRunFacet, + facet_key: str = "device", + ) -> None: + """Assert that the event contains the expected DeviceRunFacet. + + Args: + event: The emitted event + expected_facet: Expected DeviceRunFacet instance + facet_key: Key in event.run.facets (default: "device") + """ + assert facet_key in event.run.facets + facet = event.run.facets[facet_key] + assert facet.hostname == expected_facet.hostname + + def assert_aws_job_facet( + self, + event: Any, + expected_facet: AWSJobRunFacet, + facet_key: str = "awsJob", + ) -> None: + """Assert that the event contains the expected AWSJobRunFacet. + + Args: + event: The emitted event + expected_facet: Expected AWSJobRunFacet instance + facet_key: Key in event.run.facets (default: "awsJob") + """ + assert facet_key in event.run.facets + facet = event.run.facets[facet_key] + assert facet.name == expected_facet.name + assert facet.id == expected_facet.id + + def assert_nominal_time( + self, + event: Any, + expected_start: Optional[str] = None, + expected_end: Optional[str] = None, + ) -> None: + """Assert that the event contains the expected nominalTime facet. + + Args: + event: The emitted event + expected_start: Expected nominalStartTime (None = not checked) + expected_end: Expected nominalEndTime (None = not checked) + + Example: + helper.assert_nominal_time( + event, + expected_start="2025-10-23T01:00:00+00:00", + expected_end="2025-10-23T01:30:00+00:00", + ) + """ + assert "nominalTime" in event.run.facets + facet = event.run.facets["nominalTime"] + + if expected_start is not None: + assert facet.nominalStartTime == expected_start + if expected_end is not None: + assert facet.nominalEndTime == expected_end + + def assert_event_structure( + self, + event: Any, + expected_type: RunState, + expected_namespace: str, + expected_job_name: str, + expected_run_id: str, + expected_producer: str, + ) -> None: + """Assert the basic structure of an event. + + Args: + event: The emitted event + expected_type: Expected event type (START, COMPLETE, etc.) + expected_namespace: Expected job namespace + expected_job_name: Expected job name + expected_run_id: Expected run ID + expected_producer: Expected producer string + """ + assert event.eventType == expected_type + assert event.job.namespace == expected_namespace + assert event.job.name == expected_job_name + assert event.run.runId == expected_run_id + assert event.producer == expected_producer + + def assert_complete_event_structure( + self, + mock_client: MagicMock, + session: BaseSession, + expected_producer: str, + ) -> None: + """Assert the complete event structure for a completed session. + + Verifies: + - COMPLETE event is second call + - Event has correct type, namespace, job_name, run_id, producer + + Args: + mock_client: The mocked MarquezClient + session: The session instance + expected_producer: Expected producer string + """ + complete_event = self.get_emitted_event(mock_client, call_index=1) + self.assert_event_structure( + complete_event, + expected_type=RunState.COMPLETE, + expected_namespace=session.namespace, + expected_job_name=session.job_name, + expected_run_id=session.run_id, + expected_producer=expected_producer, + ) + + def __enter__(self) -> "BaseSessionTestHelper": + """Context manager entry.""" + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + """Context manager exit - automatically stop patch.""" + self.stop_patch() diff --git a/tests/unit/helpers/test_test_helper.py b/tests/unit/helpers/test_test_helper.py new file mode 100644 index 0000000..938266a --- /dev/null +++ b/tests/unit/helpers/test_test_helper.py @@ -0,0 +1,340 @@ +"""Unit tests for BaseSessionTestHelper.""" + +from typing import Dict + +from openlineage.client.facet import BaseFacet +from openlineage.client.run import RunState + +from airoa_lineage.core.base_session import BaseSession +from airoa_lineage.facets import AWSJobRunFacet, CommonRunFacet, DeviceRunFacet + +from .test_helper import BaseSessionTestHelper + + +class MinimalSession(BaseSession): + """Minimal session for testing BaseSessionTestHelper.""" + + def __init__( + self, + namespace: str, + common_facet: CommonRunFacet, + job_name: str = "test-job", + **kwargs, + ): + """Initialize minimal session.""" + super().__init__(namespace, job_name, **kwargs) + self.common_facet = common_facet + + def _get_session_facets(self) -> Dict[str, BaseFacet]: + """Return common facet.""" + return {"common": self.common_facet} + + def _get_producer(self) -> str: + """Return test producer.""" + return "test-producer" + + +class DeviceSession(BaseSession): + """Session with device facet for testing.""" + + def __init__( + self, + namespace: str, + common_facet: CommonRunFacet, + device_facet: DeviceRunFacet, + job_name: str = "test-job", + **kwargs, + ): + """Initialize session with device facet.""" + super().__init__(namespace, job_name, **kwargs) + self.common_facet = common_facet + self.device_facet = device_facet + + def _get_session_facets(self) -> Dict[str, BaseFacet]: + """Return common and device facets.""" + return {"common": self.common_facet, "device": self.device_facet} + + def _get_producer(self) -> str: + """Return test producer.""" + return "test-producer" + + +class AWSJobSession(BaseSession): + """Session with AWS job facet for testing.""" + + def __init__( + self, + namespace: str, + common_facet: CommonRunFacet, + aws_job_facet: AWSJobRunFacet, + job_name: str = "test-job", + **kwargs, + ): + """Initialize session with AWS job facet.""" + super().__init__(namespace, job_name, **kwargs) + self.common_facet = common_facet + self.aws_job_facet = aws_job_facet + + def _get_session_facets(self) -> Dict[str, BaseFacet]: + """Return common and AWS job facets.""" + return {"common": self.common_facet, "awsJob": self.aws_job_facet} + + def _get_producer(self) -> str: + """Return test producer.""" + return "test-producer" + + +class TestBaseSessionTestHelper: + """Test BaseSessionTestHelper functionality.""" + + def test_create_session(self, common_facet): + """Test that create_session creates a session with mocked client.""" + helper = BaseSessionTestHelper(MinimalSession) + + try: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + + # Verify session was created + assert isinstance(session, MinimalSession) + assert session.namespace == "test_namespace" + assert session.job_name == "test-job" + assert session.common_facet == common_facet + + # Verify mock client was created + assert mock_client is not None + finally: + helper.stop_patch() + + def test_create_session_with_custom_kwargs(self, common_facet): + """Test create_session with custom keyword arguments.""" + helper = BaseSessionTestHelper(MinimalSession) + + try: + session, mock_client = helper.create_session( + namespace="custom_namespace", + common_facet=common_facet, + job_name="custom-job", + facet_prefix="airoa", + ) + + assert session.namespace == "custom_namespace" + assert session.job_name == "custom-job" + assert session.facet_prefix == "airoa" + finally: + helper.stop_patch() + + def test_get_emitted_event(self, common_facet): + """Test that get_emitted_event extracts the correct event.""" + helper = BaseSessionTestHelper(MinimalSession) + + try: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + # Get the START event (first call) + event = helper.get_emitted_event(mock_client, call_index=0) + + assert event.eventType == RunState.START + assert event.job.namespace == "test_namespace" + assert event.job.name == "test-job" + finally: + helper.stop_patch() + + def test_assert_start_success(self, common_facet): + """Test that assert_start_success verifies start correctly.""" + helper = BaseSessionTestHelper(MinimalSession) + + try: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + # Should not raise any assertion errors + helper.assert_start_success(session, mock_client) + finally: + helper.stop_patch() + + def test_assert_complete_success(self, common_facet): + """Test that assert_complete_success verifies complete correctly.""" + helper = BaseSessionTestHelper(MinimalSession) + + try: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + session.complete() + + # Should not raise any assertion errors + helper.assert_complete_success(session, mock_client) + finally: + helper.stop_patch() + + def test_assert_common_facet(self, common_facet): + """Test that assert_common_facet verifies CommonRunFacet correctly.""" + helper = BaseSessionTestHelper(MinimalSession) + + try: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + event = helper.get_emitted_event(mock_client, call_index=0) + + # Should not raise any assertion errors + helper.assert_common_facet(event, common_facet, facet_key="common") + finally: + helper.stop_patch() + + def test_assert_common_facet_with_prefix(self, common_facet): + """Test assert_common_facet with facet prefix.""" + helper = BaseSessionTestHelper(MinimalSession) + + try: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + facet_prefix="airoa", + ) + session.start() + + event = helper.get_emitted_event(mock_client, call_index=0) + + # Should verify with prefixed key + helper.assert_common_facet(event, common_facet, facet_key="airoa_common") + finally: + helper.stop_patch() + + def test_assert_device_facet(self, common_facet, device_facet): + """Test that assert_device_facet verifies DeviceRunFacet correctly.""" + helper = BaseSessionTestHelper(DeviceSession) + + try: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + session.start() + + event = helper.get_emitted_event(mock_client, call_index=0) + + # Should not raise any assertion errors + helper.assert_device_facet(event, device_facet, facet_key="device") + finally: + helper.stop_patch() + + def test_assert_aws_job_facet(self, common_facet, aws_job_facet): + """Test that assert_aws_job_facet verifies AWSJobRunFacet correctly.""" + helper = BaseSessionTestHelper(AWSJobSession) + + try: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + aws_job_facet=aws_job_facet, + ) + session.start() + + event = helper.get_emitted_event(mock_client, call_index=0) + + # Should not raise any assertion errors + helper.assert_aws_job_facet(event, aws_job_facet, facet_key="awsJob") + finally: + helper.stop_patch() + + def test_assert_nominal_time(self, common_facet): + """Test that assert_nominal_time verifies nominal time correctly.""" + helper = BaseSessionTestHelper(MinimalSession) + + try: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start( + nominal_start_time="2025-10-23T01:00:00+00:00", + nominal_end_time="2025-10-23T01:30:00+00:00", + ) + + event = helper.get_emitted_event(mock_client, call_index=0) + + # Should not raise any assertion errors + helper.assert_nominal_time( + event, + expected_start="2025-10-23T01:00:00+00:00", + expected_end="2025-10-23T01:30:00+00:00", + ) + finally: + helper.stop_patch() + + def test_assert_event_structure(self, common_facet): + """Test that assert_event_structure verifies event structure correctly.""" + helper = BaseSessionTestHelper(MinimalSession) + + try: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + event = helper.get_emitted_event(mock_client, call_index=0) + + # Should not raise any assertion errors + helper.assert_event_structure( + event, + expected_type=RunState.START, + expected_namespace="test_namespace", + expected_job_name="test-job", + expected_run_id=session.run_id, + expected_producer="test-producer", + ) + finally: + helper.stop_patch() + + def test_assert_complete_event_structure(self, common_facet): + """Test that assert_complete_event_structure verifies complete event.""" + helper = BaseSessionTestHelper(MinimalSession) + + try: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + session.complete() + + # Should not raise any assertion errors + helper.assert_complete_event_structure( + mock_client, + session, + expected_producer="test-producer", + ) + finally: + helper.stop_patch() + + def test_context_manager(self, common_facet): + """Test that helper works as a context manager.""" + with BaseSessionTestHelper(MinimalSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + ) + session.start() + + helper.assert_start_success(session, mock_client) + + # Patch should be automatically stopped + # (Cannot easily verify without implementation details) diff --git a/tests/unit/teleop/test_session.py b/tests/unit/teleop/test_session.py index a9544aa..a512072 100644 --- a/tests/unit/teleop/test_session.py +++ b/tests/unit/teleop/test_session.py @@ -2,28 +2,19 @@ import os import uuid -from unittest.mock import MagicMock, patch import pytest from airoa_lineage.facets import CommonRunFacet, DeviceRunFacet from airoa_lineage.teleop import TeleopSession +from tests.unit.helpers.test_helper import BaseSessionTestHelper class TestTeleopSessionInitialization: """Test TeleopSession initialization.""" - def test_initialization_with_defaults(self): + def test_initialization_with_defaults(self, common_facet, device_facet): """Test initialization with default values.""" - 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="operator-pc-001") session = TeleopSession( namespace="test_namespace", common_facet=common_facet, @@ -91,452 +82,233 @@ def test_initialization_with_custom_values(self): class TestTeleopSessionStart: """Test TeleopSession start() method.""" - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_start_returns_run_id(self, mock_client_class): + def test_start_returns_run_id(self, common_facet, device_facet): """Test that start() returns run_id.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - run_id = session.start() + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + run_id = session.start() - # 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 emit was called once - mock_client.emit.assert_called_once() + # Verify run_id is returned + assert run_id == session.run_id + # Verify session state and emit + helper.assert_start_success(session, mock_client) - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_start_twice_raises_error(self, mock_client_class): + def test_start_twice_raises_error(self, common_facet, device_facet): """Test that calling start() twice raises RuntimeError.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - session.start() - - # Second call should raise error - with pytest.raises(RuntimeError, match="already started"): + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) session.start() - # Verify emit was called only once (not twice) - assert mock_client.emit.call_count == 1 + # Second call should raise error + with pytest.raises(RuntimeError, match="already started"): + session.start() + + # Verify emit was called only once (not twice) + assert mock_client.emit.call_count == 1 class TestTeleopSessionComplete: """Test TeleopSession complete() method.""" - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_complete_without_start_raises_error(self, mock_client_class): + def test_complete_without_start_raises_error(self, common_facet, device_facet): """Test that calling complete() without start() raises RuntimeError.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) - # complete() without start() should raise error - with pytest.raises(RuntimeError, match="not started"): - session.complete() + # complete() without start() should raise error + with pytest.raises(RuntimeError, match="not started"): + session.complete() - # Verify emit was never called - mock_client.emit.assert_not_called() + # Verify emit was never called + mock_client.emit.assert_not_called() - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_complete_twice_raises_error(self, mock_client_class): + def test_complete_twice_raises_error(self, common_facet, device_facet): """Test that calling complete() twice raises RuntimeError.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - session.start() - session.complete() - - # Second complete() should raise error - with pytest.raises(RuntimeError, match="already completed"): + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + session.start() session.complete() - # Verify emit was called twice (start + complete, not 3 times) - assert mock_client.emit.call_count == 2 + # Second complete() should raise error + with pytest.raises(RuntimeError, match="already completed"): + session.complete() - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_complete_success(self, mock_client_class): - """Test successful complete() after start().""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client + # Verify emit was called twice (start + complete, not 3 times) + assert mock_client.emit.call_count == 2 - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - session.start() - session.complete() + def test_complete_success(self, common_facet, device_facet): + """Test successful complete() after start().""" + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + session.start() + session.complete() - # Verify session state - assert session._started is True - assert session._completed is True - # Verify emit was called twice (start + complete) - assert mock_client.emit.call_count == 2 + # Verify session state + helper.assert_complete_success(session, mock_client) class TestTeleopSessionNominalTime: """Test TeleopSession with nominal time facets.""" - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_start_with_nominal_time_both(self, mock_client_class): + def test_start_with_nominal_time_both(self, common_facet, device_facet): """Test start() with both nominal_start_time and nominal_end_time.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - run_id = session.start( - nominal_start_time="2025-10-23T01:00:00+00:00", - nominal_end_time="2025-10-23T01:30:00+00:00", - ) + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + run_id = session.start( + nominal_start_time="2025-10-23T01:00:00+00:00", + nominal_end_time="2025-10-23T01:30:00+00:00", + ) - # Verify run_id is returned - assert run_id == session.run_id - # Verify session state changed - assert session._started is True - - # Verify emit was called with robot and nominalTime facets - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] - assert "common" in event.run.facets - assert event.run.facets["common"].robotId == "hsr001" - assert event.run.facets["common"].location == "weblab" - assert event.run.facets["common"].repositoryHash == "df110d5" - assert ( - event.run.facets["common"].repositoryUri - == "https://github.com/user/repo.git" - ) - assert event.run.facets["common"].repositoryTag == "v1.0.0" - assert event.run.facets["common"].repositoryBranch == "main" - assert "nominalTime" in event.run.facets - assert ( - event.run.facets["nominalTime"].nominalStartTime - == "2025-10-23T01:00:00+00:00" - ) - assert ( - event.run.facets["nominalTime"].nominalEndTime - == "2025-10-23T01:30:00+00:00" - ) + # Verify run_id is returned + assert run_id == session.run_id + # Verify session state changed + assert session._started is True + + # Verify emit was called with facets + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + helper.assert_common_facet(event, common_facet) + helper.assert_device_facet(event, device_facet) + helper.assert_nominal_time( + event, + expected_start="2025-10-23T01:00:00+00:00", + expected_end="2025-10-23T01:30:00+00:00", + ) - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_start_with_nominal_time_start_only(self, mock_client_class): + def test_start_with_nominal_time_start_only(self, common_facet, device_facet): """Test start() with only nominal_start_time.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - run_id = session.start(nominal_start_time="2025-10-23T01:00:00+00:00") - - # Verify run_id is returned - assert run_id == session.run_id - - # Verify emit was called with robot and nominalTime facets - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] - assert "common" in event.run.facets - assert event.run.facets["common"].robotId == "hsr001" - assert event.run.facets["common"].location == "weblab" - assert event.run.facets["common"].repositoryHash == "df110d5" - assert ( - event.run.facets["common"].repositoryUri - == "https://github.com/user/repo.git" - ) - assert event.run.facets["common"].repositoryTag == "v1.0.0" - assert event.run.facets["common"].repositoryBranch == "main" - assert "nominalTime" in event.run.facets - assert ( - event.run.facets["nominalTime"].nominalStartTime - == "2025-10-23T01:00:00+00:00" - ) - assert event.run.facets["nominalTime"].nominalEndTime is None + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + run_id = session.start(nominal_start_time="2025-10-23T01:00:00+00:00") + + # Verify run_id is returned + assert run_id == session.run_id + + # Verify emit was called with facets + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + helper.assert_common_facet(event, common_facet) + helper.assert_device_facet(event, device_facet) + helper.assert_nominal_time( + event, expected_start="2025-10-23T01:00:00+00:00" + ) + assert event.run.facets["nominalTime"].nominalEndTime is None - @patch("airoa_lineage.teleop.session.MarquezClient") - @patch("airoa_lineage.teleop.session.get_event_timestamp") - def test_start_with_nominal_time_end_only(self, mock_timestamp, mock_client_class): + def test_start_with_nominal_time_end_only(self, common_facet, device_facet): """Test start() with only nominal_end_time (start defaults to current time).""" - # Setup mocks - mock_client = MagicMock() - mock_client_class.return_value = mock_client - mock_timestamp.return_value = "2025-10-23T02:00:00+00:00" - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - run_id = session.start(nominal_end_time="2025-10-23T01:30:00+00:00") - - # Verify run_id is returned - assert run_id == session.run_id - - # Verify emit was called with robot and nominalTime facets - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] - assert "common" in event.run.facets - assert event.run.facets["common"].robotId == "hsr001" - assert event.run.facets["common"].location == "weblab" - assert event.run.facets["common"].repositoryHash == "df110d5" - assert ( - event.run.facets["common"].repositoryUri - == "https://github.com/user/repo.git" - ) - assert event.run.facets["common"].repositoryTag == "v1.0.0" - assert event.run.facets["common"].repositoryBranch == "main" - assert "nominalTime" in event.run.facets - # nominal_start_time should default to current time - assert ( - event.run.facets["nominalTime"].nominalStartTime - == "2025-10-23T02:00:00+00:00" - ) - assert ( - event.run.facets["nominalTime"].nominalEndTime - == "2025-10-23T01:30:00+00:00" - ) + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + run_id = session.start(nominal_end_time="2025-10-23T01:30:00+00:00") + + # Verify run_id is returned + assert run_id == session.run_id + + # Verify emit was called with facets + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + helper.assert_common_facet(event, common_facet) + helper.assert_device_facet(event, device_facet) + assert "nominalTime" in event.run.facets + # nominal_start_time should default to current time + assert event.run.facets["nominalTime"].nominalStartTime is not None + helper.assert_nominal_time(event, expected_end="2025-10-23T01:30:00+00:00") + + def test_start_without_nominal_time(self, common_facet, device_facet): + """Test start() without nominal_time but with facets.""" + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + run_id = session.start() - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_start_without_nominal_time(self, mock_client_class): - """Test start() without nominal_time but with robot facet.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client + # Verify run_id is returned + assert run_id == session.run_id - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - run_id = session.start() - - # Verify run_id is returned - assert run_id == session.run_id - - # Verify emit was called with robot facet but without nominalTime facet - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] - assert "common" in event.run.facets - assert event.run.facets["common"].robotId == "hsr001" - assert event.run.facets["common"].location == "weblab" - assert event.run.facets["common"].repositoryHash == "df110d5" - assert ( - event.run.facets["common"].repositoryUri - == "https://github.com/user/repo.git" - ) - assert event.run.facets["common"].repositoryTag == "v1.0.0" - assert event.run.facets["common"].repositoryBranch == "main" - assert "nominalTime" not in event.run.facets + # Verify emit was called with facets but without nominalTime facet + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) + helper.assert_common_facet(event, common_facet) + helper.assert_device_facet(event, device_facet) + assert "nominalTime" not in event.run.facets class TestCommonFacet: """Test common facet functionality.""" - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_start_includes_common_facet(self, mock_client_class): + def test_start_includes_common_facet(self, common_facet, device_facet): """Test that START event includes common facet.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - session.start() + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + session.start() - # Verify emit was called - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) - # Verify robot facet exists and has correct value - assert "common" in event.run.facets - assert event.run.facets["common"].robotId == "hsr001" - assert event.run.facets["common"].location == "weblab" - assert event.run.facets["common"].repositoryHash == "df110d5" - assert ( - event.run.facets["common"].repositoryUri - == "https://github.com/user/repo.git" - ) - assert event.run.facets["common"].repositoryTag == "v1.0.0" - assert event.run.facets["common"].repositoryBranch == "main" + # Verify common facet exists and has correct value + helper.assert_common_facet(event, common_facet) - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_complete_includes_common_facet(self, mock_client_class): + def test_complete_includes_common_facet(self, common_facet, device_facet): """Test that COMPLETE event includes common facet.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - session.start() - session.complete() - - # Verify emit was called twice - assert mock_client.emit.call_count == 2 + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + session.start() + session.complete() - # Check COMPLETE event (second call) - complete_event = mock_client.emit.call_args_list[1][0][0] + # Verify emit was called twice + assert mock_client.emit.call_count == 2 - # Verify common facet exists in COMPLETE event - assert "common" in complete_event.run.facets - assert complete_event.run.facets["common"].robotId == "hsr001" - assert complete_event.run.facets["common"].location == "weblab" - assert complete_event.run.facets["common"].repositoryHash == "df110d5" - assert ( - complete_event.run.facets["common"].repositoryUri - == "https://github.com/user/repo.git" - ) - assert complete_event.run.facets["common"].repositoryTag == "v1.0.0" - assert complete_event.run.facets["common"].repositoryBranch == "main" + # Check COMPLETE event (second call) + complete_event = helper.get_emitted_event(mock_client, call_index=1) - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_different_robot_ids_and_locations(self, mock_client_class): - """Test that different robot_ids, locations, and repository info are correctly stored.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client + # Verify common facet exists in COMPLETE event + helper.assert_common_facet(complete_event, common_facet) - # Test multiple combinations - test_cases = [ + @pytest.mark.parametrize( + "robot_id,location,repo_hash,repo_uri,repo_tag,repo_branch", + [ ( "hsr001", "weblab", @@ -569,16 +341,20 @@ def test_different_robot_ids_and_locations(self, mock_client_class): "v4.0.0", "release", ), - ] - - for ( - robot_id, - location, - repo_hash, - repo_uri, - repo_tag, - repo_branch, - ) in test_cases: + ], + ) + def test_different_robot_ids_and_locations( + self, + device_facet, + robot_id, + location, + repo_hash, + repo_uri, + repo_tag, + repo_branch, + ): + """Test that different robot_ids, locations, and repository info are correctly stored.""" + with BaseSessionTestHelper(TeleopSession) as helper: common_facet = CommonRunFacet( robotId=robot_id, location=location, @@ -587,8 +363,7 @@ def test_different_robot_ids_and_locations(self, mock_client_class): repositoryTag=repo_tag, repositoryBranch=repo_branch, ) - device_facet = DeviceRunFacet(hostname="operator-pc-001") - session = TeleopSession( + session, mock_client = helper.create_session( namespace="test_namespace", common_facet=common_facet, device_facet=device_facet, @@ -596,112 +371,64 @@ def test_different_robot_ids_and_locations(self, mock_client_class): session.start() # Verify the common facet in the event - event = mock_client.emit.call_args[0][0] - assert event.run.facets["common"].robotId == robot_id - assert event.run.facets["common"].location == location - assert event.run.facets["common"].repositoryHash == repo_hash - assert event.run.facets["common"].repositoryUri == repo_uri - assert event.run.facets["common"].repositoryTag == repo_tag - assert event.run.facets["common"].repositoryBranch == repo_branch - - # Reset mock for next iteration - mock_client.reset_mock() + event = helper.get_emitted_event(mock_client) + helper.assert_common_facet(event, common_facet) class TestDeviceFacet: """Test device facet functionality.""" - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_start_includes_device_facet(self, mock_client_class): + def test_start_includes_device_facet(self, common_facet, device_facet): """Test that START event includes device facet.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - session.start() + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + session.start() - # Verify emit was called - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) - # Verify device facet exists and has correct value - assert "device" in event.run.facets - assert event.run.facets["device"].hostname == "operator-pc-001" + # Verify device facet exists and has correct value + helper.assert_device_facet(event, device_facet) - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_complete_includes_device_facet(self, mock_client_class): + def test_complete_includes_device_facet(self, common_facet, device_facet): """Test that COMPLETE event includes device facet.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - session.start() - session.complete() - - # Verify emit was called twice - assert mock_client.emit.call_count == 2 + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + session.start() + session.complete() - # Check COMPLETE event (second call) - complete_event = mock_client.emit.call_args_list[1][0][0] + # Verify emit was called twice + assert mock_client.emit.call_count == 2 - # Verify device facet exists in COMPLETE event - assert "device" in complete_event.run.facets - assert complete_event.run.facets["device"].hostname == "operator-pc-001" + # Check COMPLETE event (second call) + complete_event = helper.get_emitted_event(mock_client, call_index=1) - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_different_hostnames(self, mock_client_class): - """Test that different hostnames are correctly stored.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client + # Verify device facet exists in COMPLETE event + helper.assert_device_facet(complete_event, device_facet) - # Test multiple hostnames - test_hostnames = [ + @pytest.mark.parametrize( + "hostname", + [ "operator-pc-001", "lab-machine-42", "dev-workstation", "research-laptop", - ] - - for hostname in test_hostnames: - common_facet = CommonRunFacet( - robotId="hsr001", - location="weblab", - repositoryHash="df110d5", - repositoryUri="https://github.com/user/repo.git", - repositoryTag="v1.0.0", - repositoryBranch="main", - ) + ], + ) + def test_different_hostnames(self, common_facet, hostname): + """Test that different hostnames are correctly stored.""" + with BaseSessionTestHelper(TeleopSession) as helper: device_facet = DeviceRunFacet(hostname=hostname) - session = TeleopSession( + session, mock_client = helper.create_session( namespace="test_namespace", common_facet=common_facet, device_facet=device_facet, @@ -709,31 +436,15 @@ def test_different_hostnames(self, mock_client_class): session.start() # Verify the device facet in the event - event = mock_client.emit.call_args[0][0] - assert event.run.facets["device"].hostname == hostname - - # Reset mock for next iteration - mock_client.reset_mock() + event = helper.get_emitted_event(mock_client) + helper.assert_device_facet(event, device_facet) class TestFacetPrefix: """Test facet_prefix functionality.""" - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_default_facet_prefix_empty(self, mock_client_class): + def test_default_facet_prefix_empty(self, common_facet, device_facet): """Test that default facet_prefix is empty string.""" - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") session = TeleopSession( namespace="test_namespace", common_facet=common_facet, @@ -743,170 +454,121 @@ def test_default_facet_prefix_empty(self, mock_client_class): # Verify facet_prefix defaults to empty string assert session.facet_prefix == "" - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_start_without_prefix_uses_default_keys(self, mock_client_class): + def test_start_without_prefix_uses_default_keys(self, common_facet, device_facet): """Test that START event uses 'common' and 'device' keys without prefix.""" - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - session.start() + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + session.start() - # Verify emit was called - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) - # Verify facet keys are default (no prefix) - assert "common" in event.run.facets - assert "device" in event.run.facets - assert event.run.facets["common"].robotId == "hsr001" - assert event.run.facets["device"].hostname == "operator-pc-001" + # Verify facet keys are default (no prefix) + assert "common" in event.run.facets + assert "device" in event.run.facets + helper.assert_common_facet(event, common_facet) + helper.assert_device_facet(event, device_facet) - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_start_with_prefix_uses_prefixed_keys(self, mock_client_class): + def test_start_with_prefix_uses_prefixed_keys(self, common_facet, device_facet): """Test that START event uses prefixed keys when facet_prefix is set.""" - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - facet_prefix="airoa", - ) - session.start() + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + facet_prefix="airoa", + ) + session.start() - # Verify emit was called - mock_client.emit.assert_called_once() - event = mock_client.emit.call_args[0][0] + # Verify emit was called + mock_client.emit.assert_called_once() + event = helper.get_emitted_event(mock_client) - # Verify facet keys have prefix - assert "airoa_common" in event.run.facets - assert "airoa_device" in event.run.facets - assert event.run.facets["airoa_common"].robotId == "hsr001" - assert event.run.facets["airoa_device"].hostname == "operator-pc-001" + # Verify facet keys have prefix + assert "airoa_common" in event.run.facets + assert "airoa_device" in event.run.facets + helper.assert_common_facet(event, common_facet, facet_key="airoa_common") + helper.assert_device_facet(event, device_facet, facet_key="airoa_device") - # Verify unprefixed keys do not exist - assert "common" not in event.run.facets - assert "device" not in event.run.facets + # Verify unprefixed keys do not exist + assert "common" not in event.run.facets + assert "device" not in event.run.facets - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_complete_without_prefix_uses_default_keys(self, mock_client_class): + def test_complete_without_prefix_uses_default_keys( + self, common_facet, device_facet + ): """Test that COMPLETE event uses 'common' and 'device' keys without prefix.""" - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - ) - session.start() - session.complete() + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + ) + session.start() + session.complete() - # Verify emit was called twice - assert mock_client.emit.call_count == 2 + # Verify emit was called twice + assert mock_client.emit.call_count == 2 - # Check COMPLETE event (second call) - complete_event = mock_client.emit.call_args_list[1][0][0] + # Check COMPLETE event (second call) + complete_event = helper.get_emitted_event(mock_client, call_index=1) - # Verify facet keys are default (no prefix) - assert "common" in complete_event.run.facets - assert "device" in complete_event.run.facets - assert complete_event.run.facets["common"].robotId == "hsr001" - assert complete_event.run.facets["device"].hostname == "operator-pc-001" + # Verify facet keys are default (no prefix) + assert "common" in complete_event.run.facets + assert "device" in complete_event.run.facets + helper.assert_common_facet(complete_event, common_facet) + helper.assert_device_facet(complete_event, device_facet) - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_complete_with_prefix_uses_prefixed_keys(self, mock_client_class): + def test_complete_with_prefix_uses_prefixed_keys(self, common_facet, device_facet): """Test that COMPLETE event uses prefixed keys when facet_prefix is set.""" - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - 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="operator-pc-001") - session = TeleopSession( - namespace="test_namespace", - common_facet=common_facet, - device_facet=device_facet, - facet_prefix="airoa", - ) - session.start() - session.complete() - - # Verify emit was called twice - assert mock_client.emit.call_count == 2 + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( + namespace="test_namespace", + common_facet=common_facet, + device_facet=device_facet, + facet_prefix="airoa", + ) + session.start() + session.complete() - # Check COMPLETE event (second call) - complete_event = mock_client.emit.call_args_list[1][0][0] + # Verify emit was called twice + assert mock_client.emit.call_count == 2 - # Verify facet keys have prefix - assert "airoa_common" in complete_event.run.facets - assert "airoa_device" in complete_event.run.facets - assert complete_event.run.facets["airoa_common"].robotId == "hsr001" - assert complete_event.run.facets["airoa_device"].hostname == "operator-pc-001" + # Check COMPLETE event (second call) + complete_event = helper.get_emitted_event(mock_client, call_index=1) - # Verify unprefixed keys do not exist - assert "common" not in complete_event.run.facets - assert "device" not in complete_event.run.facets + # Verify facet keys have prefix + assert "airoa_common" in complete_event.run.facets + assert "airoa_device" in complete_event.run.facets + helper.assert_common_facet( + complete_event, common_facet, facet_key="airoa_common" + ) + helper.assert_device_facet( + complete_event, device_facet, facet_key="airoa_device" + ) - @patch("airoa_lineage.teleop.session.MarquezClient") - def test_different_prefixes(self, mock_client_class): + # Verify unprefixed keys do not exist + assert "common" not in complete_event.run.facets + assert "device" not in complete_event.run.facets + + @pytest.mark.parametrize( + "prefix", + [ + "airoa", + "custom", + "my_prefix", + "test123", + ], + ) + def test_different_prefixes(self, common_facet, device_facet, prefix): """Test that different prefixes are correctly applied.""" - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - # Test multiple prefixes - test_prefixes = ["airoa", "custom", "my_prefix", "test123"] - - for prefix in test_prefixes: - 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="operator-pc-001") - session = TeleopSession( + with BaseSessionTestHelper(TeleopSession) as helper: + session, mock_client = helper.create_session( namespace="test_namespace", common_facet=common_facet, device_facet=device_facet, @@ -915,14 +577,15 @@ def test_different_prefixes(self, mock_client_class): session.start() # Verify the facet keys in the event - event = mock_client.emit.call_args[0][0] + event = helper.get_emitted_event(mock_client) expected_common_key = f"{prefix}_common" expected_device_key = f"{prefix}_device" assert expected_common_key in event.run.facets assert expected_device_key in event.run.facets - assert event.run.facets[expected_common_key].robotId == "hsr001" - assert event.run.facets[expected_device_key].hostname == "operator-pc-001" - - # Reset mock for next iteration - mock_client.reset_mock() + helper.assert_common_facet( + event, common_facet, facet_key=expected_common_key + ) + helper.assert_device_facet( + event, device_facet, facet_key=expected_device_key + ) diff --git a/tests/unit/usb_copy/test_session.py b/tests/unit/usb_copy/test_session.py index ded79ba..434917d 100644 --- a/tests/unit/usb_copy/test_session.py +++ b/tests/unit/usb_copy/test_session.py @@ -1,55 +1,19 @@ """Unit tests for USBCopySession.""" -import os import uuid -from unittest.mock import MagicMock, patch -import pytest from airoa_lineage.facets import CommonRunFacet from airoa_lineage.usb_copy import USBCopySession +from tests.unit.helpers.base_session_test import BaseSessionTest -@pytest.fixture -def common_facet(): - """Fixture for common run facet.""" - return CommonRunFacet( - robotId="hsr001", - location="weblab", - repositoryHash="df110d5", - repositoryUri="https://github.com/user/repo.git", - repositoryTag="v1.0.0", - repositoryBranch="main", - ) - - -class TestUSBCopySessionInitialization: +class TestUSBCopySessionInitialization(BaseSessionTest): """Test USBCopySession initialization.""" - def test_initialization_with_defaults(self, common_facet): - """Test initialization with default values.""" - session = USBCopySession( - namespace="test_namespace", - common_facet=common_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.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 + SESSION_CLASS = USBCopySession + DEFAULT_JOB_NAME = "usb-data-copy" + PRODUCER_NAME = "airoa-usbcopy-system" def test_initialization_with_custom_values(self): """Test initialization with custom values.""" @@ -85,380 +49,42 @@ def test_initialization_with_custom_values(self): assert session._started is False assert session._completed is False - def test_initialization_with_facet_prefix(self, common_facet): - """Test initialization with facet_prefix.""" - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - facet_prefix="airoa", - ) - - assert session.facet_prefix == "airoa" - -class TestUSBCopySessionStart: +class TestUSBCopySessionStart(BaseSessionTest): """Test USBCopySession start() method.""" - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_start_returns_run_id(self, mock_client_class, common_facet): - """Test that start() returns run_id.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - run_id = session.start() - - # 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 emit was called once - mock_client.emit.assert_called_once() - - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_start_twice_raises_error(self, mock_client_class, common_facet): - """Test that calling start() twice raises RuntimeError.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - session.start() - - # Second call should raise error - with pytest.raises(RuntimeError, match="already started"): - session.start() - - # Verify emit was called only once (not twice) - assert mock_client.emit.call_count == 1 - - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_start_event_has_correct_producer(self, mock_client_class, common_facet): - """Test that START event has correct producer.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - session.start() + SESSION_CLASS = USBCopySession + DEFAULT_JOB_NAME = "usb-data-copy" + PRODUCER_NAME = "airoa-usbcopy-system" - # Get the event that was passed to emit() - event = mock_client.emit.call_args[0][0] - assert event.producer == "airoa-usbcopy-system" - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_start_event_has_empty_inputs_outputs( - self, mock_client_class, common_facet - ): - """Test that START event has empty inputs and outputs.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - session.start() - - # Get the event that was passed to emit() - event = mock_client.emit.call_args[0][0] - assert event.inputs == [] - assert event.outputs == [] - - -class TestUSBCopySessionComplete: +class TestUSBCopySessionComplete(BaseSessionTest): """Test USBCopySession complete() method.""" - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_complete_without_start_raises_error(self, mock_client_class, common_facet): - """Test that calling complete() without start() raises RuntimeError.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - - # complete() without start() should raise error - with pytest.raises(RuntimeError, match="not started"): - session.complete() - - # Verify emit was never called - mock_client.emit.assert_not_called() - - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_complete_twice_raises_error(self, mock_client_class, common_facet): - """Test that calling complete() twice raises RuntimeError.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - session.start() - session.complete() - - # Second call should raise error - with pytest.raises(RuntimeError, match="already completed"): - session.complete() + SESSION_CLASS = USBCopySession + DEFAULT_JOB_NAME = "usb-data-copy" + PRODUCER_NAME = "airoa-usbcopy-system" - # Verify emit was called twice (start + complete once) - assert mock_client.emit.call_count == 2 - - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_complete_sets_completed_flag(self, mock_client_class, common_facet): - """Test that complete() sets _completed flag.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - session.start() - session.complete() - - # Verify session state - assert session._started is True - assert session._completed is True - - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_complete_event_has_correct_producer(self, mock_client_class, common_facet): - """Test that COMPLETE event has correct producer.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - session.start() - session.complete() - # Get the COMPLETE event (second call) - event = mock_client.emit.call_args_list[1][0][0] - assert event.producer == "airoa-usbcopy-system" - - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_complete_event_has_empty_inputs_outputs( - self, mock_client_class, common_facet - ): - """Test that COMPLETE event has empty inputs and outputs.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - session.start() - session.complete() - - # Get the COMPLETE event (second call) - event = mock_client.emit.call_args_list[1][0][0] - assert event.inputs == [] - assert event.outputs == [] - - -class TestUSBCopySessionNominalTime: +class TestUSBCopySessionNominalTime(BaseSessionTest): """Test USBCopySession nominal time support.""" - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_start_with_both_nominal_times(self, mock_client_class, common_facet): - """Test start() with both nominal start and end times.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - session.start( - nominal_start_time="2025-10-23T01:00:00+00:00", - nominal_end_time="2025-10-23T01:30:00+00:00", - ) - - # Get the event that was passed to emit() - event = mock_client.emit.call_args[0][0] - # Verify nominalTime facet exists - assert "nominalTime" in event.run.facets - assert ( - event.run.facets["nominalTime"].nominalStartTime - == "2025-10-23T01:00:00+00:00" - ) - assert ( - event.run.facets["nominalTime"].nominalEndTime - == "2025-10-23T01:30:00+00:00" - ) - - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_start_with_only_nominal_end_time(self, mock_client_class, common_facet): - """Test start() with only nominal end time.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - session.start(nominal_end_time="2025-10-23T01:30:00+00:00") - - # Get the event that was passed to emit() - event = mock_client.emit.call_args[0][0] - # Verify nominalTime facet exists - assert "nominalTime" in event.run.facets - # nominalStartTime should be auto-filled with current time - assert event.run.facets["nominalTime"].nominalStartTime is not None - assert ( - event.run.facets["nominalTime"].nominalEndTime - == "2025-10-23T01:30:00+00:00" - ) - - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_start_without_nominal_time(self, mock_client_class, common_facet): - """Test start() without nominal time.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - session.start() - - # Get the event that was passed to emit() - event = mock_client.emit.call_args[0][0] - # Verify nominalTime facet does not exist - assert "nominalTime" not in event.run.facets + SESSION_CLASS = USBCopySession + DEFAULT_JOB_NAME = "usb-data-copy" + PRODUCER_NAME = "airoa-usbcopy-system" -class TestCommonFacet: +class TestCommonFacet(BaseSessionTest): """Test CommonRunFacet in USBCopySession.""" - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_start_includes_common_facet(self, mock_client_class, common_facet): - """Test that START event includes common facet.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - session.start() - - # Get the event that was passed to emit() - event = mock_client.emit.call_args[0][0] - # Verify common facet exists - assert "common" in event.run.facets - assert event.run.facets["common"].robotId == "hsr001" - assert event.run.facets["common"].location == "weblab" + SESSION_CLASS = USBCopySession + DEFAULT_JOB_NAME = "usb-data-copy" + PRODUCER_NAME = "airoa-usbcopy-system" - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_complete_includes_common_facet(self, mock_client_class, common_facet): - """Test that COMPLETE event includes common facet.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - ) - session.start() - session.complete() - - # Get the COMPLETE event (second call) - event = mock_client.emit.call_args_list[1][0][0] - # Verify common facet exists - assert "common" in event.run.facets - assert event.run.facets["common"].robotId == "hsr001" - assert event.run.facets["common"].location == "weblab" - - -class TestFacetPrefix: +class TestFacetPrefix(BaseSessionTest): """Test facet_prefix functionality.""" - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_facet_prefix_in_start_event(self, mock_client_class, common_facet): - """Test that facet_prefix is applied in START event.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - facet_prefix="airoa", - ) - session.start() - - # Get the event that was passed to emit() - event = mock_client.emit.call_args[0][0] - # Verify common facet has prefix - assert "airoa_common" in event.run.facets - assert event.run.facets["airoa_common"].robotId == "hsr001" - - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_facet_prefix_in_complete_event(self, mock_client_class, common_facet): - """Test that facet_prefix is applied in COMPLETE event.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - facet_prefix="airoa", - ) - session.start() - session.complete() - - # Get the COMPLETE event (second call) - event = mock_client.emit.call_args_list[1][0][0] - # Verify common facet has prefix - assert "airoa_common" in event.run.facets - assert event.run.facets["airoa_common"].robotId == "hsr001" - - @patch("airoa_lineage.usb_copy.session.MarquezClient") - def test_no_prefix_uses_default_key(self, mock_client_class, common_facet): - """Test that no prefix uses default facet key.""" - # Setup mock - mock_client = MagicMock() - mock_client_class.return_value = mock_client - - session = USBCopySession( - namespace="test_namespace", - common_facet=common_facet, - facet_prefix="", - ) - session.start() - - # Get the event that was passed to emit() - event = mock_client.emit.call_args[0][0] - # Verify common facet uses default key - assert "common" in event.run.facets - assert "airoa_common" not in event.run.facets + SESSION_CLASS = USBCopySession + DEFAULT_JOB_NAME = "usb-data-copy" + PRODUCER_NAME = "airoa-usbcopy-system" diff --git a/tests/unit/wasabi_upload/__init__.py b/tests/unit/wasabi_upload/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/wasabi_upload/test_session.py b/tests/unit/wasabi_upload/test_session.py new file mode 100644 index 0000000..ff95a62 --- /dev/null +++ b/tests/unit/wasabi_upload/test_session.py @@ -0,0 +1,90 @@ +"""Unit tests for WasabiUploadSession.""" + +import uuid + + +from airoa_lineage.facets import CommonRunFacet +from airoa_lineage.wasabi_upload import WasabiUploadSession +from tests.unit.helpers.base_session_test import BaseSessionTest + + +class TestWasabiUploadSessionInitialization(BaseSessionTest): + """Test WasabiUploadSession initialization.""" + + SESSION_CLASS = WasabiUploadSession + DEFAULT_JOB_NAME = "wasabi-data-upload" + PRODUCER_NAME = "airoa-wasabi-system" + + def test_initialization_with_custom_values(self): + """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 = WasabiUploadSession( + namespace="custom_namespace", + common_facet=common_facet, + job_name="custom_job", + 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.marquez_url == "http://example.com:9000" + assert session.run_id == custom_run_id + assert session._started is False + assert session._completed is False + + +class TestWasabiUploadSessionStart(BaseSessionTest): + """Test WasabiUploadSession start() method.""" + + SESSION_CLASS = WasabiUploadSession + DEFAULT_JOB_NAME = "wasabi-data-upload" + PRODUCER_NAME = "airoa-wasabi-system" + + +class TestWasabiUploadSessionComplete(BaseSessionTest): + """Test WasabiUploadSession complete() method.""" + + SESSION_CLASS = WasabiUploadSession + DEFAULT_JOB_NAME = "wasabi-data-upload" + PRODUCER_NAME = "airoa-wasabi-system" + + +class TestWasabiUploadSessionNominalTime(BaseSessionTest): + """Test WasabiUploadSession nominal time support.""" + + SESSION_CLASS = WasabiUploadSession + DEFAULT_JOB_NAME = "wasabi-data-upload" + PRODUCER_NAME = "airoa-wasabi-system" + + +class TestCommonFacet(BaseSessionTest): + """Test CommonRunFacet in WasabiUploadSession.""" + + SESSION_CLASS = WasabiUploadSession + DEFAULT_JOB_NAME = "wasabi-data-upload" + PRODUCER_NAME = "airoa-wasabi-system" + + +class TestFacetPrefix(BaseSessionTest): + """Test facet_prefix functionality.""" + + SESSION_CLASS = WasabiUploadSession + DEFAULT_JOB_NAME = "wasabi-data-upload" + PRODUCER_NAME = "airoa-wasabi-system"