diff --git a/.gitignore b/.gitignore index 0c9a791..b9198c7 100644 --- a/.gitignore +++ b/.gitignore @@ -37,7 +37,10 @@ __pycache__/ # Ignorar coverage htmlcov/ .coverage +coverage.xml git_log -.env \ No newline at end of file +.env +.ruff_cache/ +.mypy_cache/ \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..517fead --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,159 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "scouter" +version = "0.0.0" +description = "Sientia DataOps Scouter - ML Model Orchestration System" +readme = "README.md" +requires-python = ">=3.11" +authors = [ + {name = "Aignosi", email = "dev@aignosi.com"} +] + +[tool.ruff] +line-length = 100 +target-version = "py311" +exclude = [ + ".git", + ".venv", + "venv", + "__pycache__", + "*.pyc", + ".pytest_cache", + "htmlcov", +] + +[tool.ruff.lint] +select = [ + "E", # pycodestyle errors + "W", # pycodestyle warnings + "F", # pyflakes + "I", # isort + "B", # flake8-bugbear + "C4", # flake8-comprehensions + "UP", # pyupgrade + "N", # pep8-naming + "YTT", # flake8-2020 + "S", # flake8-bandit + "BLE", # flake8-blind-except + "A", # flake8-builtins + "C90", # mccabe complexity +] + +ignore = [ + "BLE001", # ignore blind except, we need to send notifications with any error + "E501", # line too long (handled by formatter) + "S101", # use of assert (needed for tests) + "S105", # possible hardcoded password (false positives) + "S106", # possible hardcoded password (false positives) + "N802", # function name should be lowercase (temporal decorators) + "N806", # variable in function should be lowercase +] + +[tool.ruff.lint.per-file-ignores] +"tests/**/*.py" = [ + "S101", # assert allowed in tests + "S105", # hardcoded passwords ok in tests + "S106", # hardcoded passwords ok in tests +] + +[tool.ruff.lint.mccabe] +max-complexity = 15 + +[tool.ruff.format] +quote-style = "single" +indent-style = "space" +line-ending = "auto" + +[tool.mypy] +python_version = "3.11" +warn_return_any = false +warn_unused_configs = true +disallow_untyped_defs = false +disallow_incomplete_defs = false +check_untyped_defs = true +no_implicit_optional = true +warn_redundant_casts = true +warn_unused_ignores = false +warn_no_return = true +strict_equality = true +ignore_missing_imports = true + +# Ignore missing imports for external packages +[[tool.mypy.overrides]] +module = "temporalio.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "sientia_do.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "mlflow.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "prometheus_client.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "sientia.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "pandas.*" +ignore_missing_imports = true + +[tool.pytest.ini_options] +testpaths = ["tests"] +python_files = ["test_*.py"] +python_classes = ["Test*"] +python_functions = ["test_*"] +addopts = [ + "-v", + "--strict-markers", + "--cov=model_manager", + "--cov-report=term-missing", + "--cov-report=html", + "--cov-report=xml", +] +markers = [ + "asyncio: marks tests as async", + "integration: marks tests as integration tests", + "unit: marks tests as unit tests", +] + +[tool.coverage.run] +source = ["model_manager"] +omit = [ + "*/tests/*", + "*/venv/*", + "*/__pycache__/*", + "*/site-packages/*", +] +branch = true + +[tool.coverage.report] +precision = 2 +show_missing = true +skip_covered = false +exclude_lines = [ + "pragma: no cover", + "def __repr__", + "def __str__", + "raise AssertionError", + "raise NotImplementedError", + "if __name__ == .__main__.:", + "if TYPE_CHECKING:", + "class .*\\bProtocol\\):", + "@(abc\\.)?abstractmethod", +] + +[tool.coverage.html] +directory = "htmlcov" + +[tool.bandit] +exclude_dirs = ["tests", "venv", ".venv"] +skips = ["B101", "B601"] # Skip assert and shell injection in controlled environments \ No newline at end of file diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..56ab376 --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,19 @@ +# Development and Testing Dependencies +# These packages are only needed for development, testing, and code quality checks +# Install with: pip install -r requirements-dev.txt + +# Code Quality & Linting +ruff>=0.1.0 # Fast Python linter and formatter (replaces flake8, black, isort) +mypy>=1.7.0 # Static type checker +bandit>=1.7.5 # Security vulnerability scanner +pandas-stubs>=2.0.0 # Type stubs for pandas +types-requests>=2.31.0 # Type stubs for requests + +# Testing +pytest>=7.4.0 # Testing framework +pytest-cov>=4.1.0 # Coverage plugin for pytest +pytest-asyncio>=0.21.0 # Async test support (already in main requirements) + +# Development Tools +ipython>=8.12.0 # Enhanced Python shell +ipdb>=0.13.13 # IPython debugger \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 8c7edc9..3d75d5e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,7 +5,6 @@ asyncua redis aiokafka pymongo -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.4.3 -git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.38.5 +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.4.6 pydruid[pandas] prometheus-client \ No newline at end of file diff --git a/scouter/activities/activities.py b/scouter/activities/activities.py index 86babf5..79ddfe1 100644 --- a/scouter/activities/activities.py +++ b/scouter/activities/activities.py @@ -1,14 +1,16 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from sientia_do.temporal.activities.postgres import Postgres + from os import getenv + from typing import Any + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.observability.logger import Logger - from scouter.activities.redis import Redis + from sientia_do.temporal.activities.postgres import Postgres + from scouter.activities.gates import Gates from scouter.activities.mongodb import MongoDB - from typing import Any - from os import getenv + from scouter.activities.redis import Redis class Activities(Postgres, Redis, Gates, MongoDB): @@ -27,12 +29,14 @@ class Activities(Postgres, Redis, Gates, MongoDB): interface while maintaining separation of concerns across different data services. """ - def __init__(self, - postgres_config: dict[str, Any], - redis_config: dict[str, Any], - mongodb_config: dict[str, Any], - logger: Logger, - notification_handler: NotificationHandler): + def __init__( + self, + postgres_config: dict[str, Any], + redis_config: dict[str, Any], + mongodb_config: dict[str, Any], + logger: Logger, + notification_handler: NotificationHandler, + ): """ Initialize the Activities class with all required services. @@ -57,7 +61,7 @@ class Activities(Postgres, Redis, Gates, MongoDB): min_connections=postgres_config['min_connections'], max_connections=postgres_config['max_connections'], logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) # Initialize Redis @@ -68,15 +72,11 @@ class Activities(Postgres, Redis, Gates, MongoDB): logger=logger, notification_handler=notification_handler, username=redis_config['username'], - password=redis_config['password'] + password=redis_config['password'], ) # Initialize Gates - Gates.__init__( - self, - logger=logger, - notification_handler=notification_handler - ) + Gates.__init__(self, logger=logger, notification_handler=notification_handler) # Initialize MongoDB MongoDB.__init__( @@ -84,10 +84,10 @@ class Activities(Postgres, Redis, Gates, MongoDB): connection_string=mongodb_config['connection_string'], database_name=mongodb_config['database_name'], logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) - self.pod_id = getenv("HOSTNAME", "localhost") + self.pod_id = getenv('HOSTNAME', 'localhost') def shutdown(self): """ diff --git a/scouter/activities/faker.py b/scouter/activities/faker.py deleted file mode 100644 index 25e3f18..0000000 --- a/scouter/activities/faker.py +++ /dev/null @@ -1,117 +0,0 @@ -import random -from datetime import datetime, timezone -from typing import Any -import json -from kafka import KafkaProducer -from temporalio import activity - -from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler -from sientia_do.temporal.activities.base import BaseActivity -from sientia_do.observability.logger import Logger - - -class Faker(BaseActivity): - """ - Synthetic data generation for testing and development. - - This class generates realistic industrial sensor data for testing purposes. - It provides: - - Configurable sensor tag simulation - - Realistic data value generation - - Kafka integration for data publishing - - Comprehensive error handling and logging - - The class is designed for development, testing, and demonstration of - data processing pipelines without requiring real industrial data sources. - """ - - def __init__(self, bootstrap_servers: str, logger: Logger, - notification_handler: NotificationHandler): - """ - Initialize the Faker class with Kafka producer and sensor configuration. - - Args: - bootstrap_servers (str): Kafka bootstrap servers configuration - logger (Logger): Logger instance for operation logging - notification_handler (NotificationHandler): Handler for system notifications - """ - self.producer = KafkaProducer( - bootstrap_servers=bootstrap_servers, - value_serializer=lambda v: json.dumps(v).encode('utf-8') - ) - - # Predefined lists for tag and name - self.tags = { - 'ns=1;i=1001': 'Temperature Sensor', - 'ns=1;i=1002': 'Vibration Meter', - 'ns=1;i=1003': 'Pressure Gauge', - 'ns=1;i=1004': 'Flow Meter', - 'ns=1;i=1005': 'Voltage Sensor', - 'ns=1;i=1006': 'Current Sensor' - } - - BaseActivity.__init__(self, logger, notification_handler) - - @activity.defn(name="generate_and_send_data") - async def generate_and_send_data(self, input_data: dict[str, Any]) -> None: - """ - Generate synthetic sensor data and publish to Kafka topic. - - This activity creates realistic industrial sensor readings and publishes - them to the specified Kafka topic. The data includes sensor tags, names, - timestamps, and values with configurable message counts. - - Args: - input_data (dict[str, Any]): Activity input parameters. - Required fields: - - topic (str): Kafka topic name for data publication - - metadata (dict[str, Any], optional): Workflow execution metadata - - num_messages (int, optional): Number of messages to generate. - Defaults to random count between 1 and available sensor tags - - Returns: - None: This activity publishes data but doesn't return results - - Raises: - ValueError: If topic is not specified - Exception: If data generation or Kafka publishing fails - """ - metadata = input_data['metadata'] - topic = input_data.get('topic') - num_messages = input_data.get( - 'num_messages', random.randint(1, len(self.tags))) # NOSONAR - - if not topic: - raise ValueError("Topic must be specified in input_data") - - self.info( - f"Generating {num_messages} messages for topic {topic}", - metadata=metadata - ) - - for _ in range(num_messages): - # Select random tag and name - tag = random.choice(list(self.tags.keys())) # NOSONAR - name = self.tags[tag] - - # Generate random value between 0 and 100 - if random.random() < 0.1: # NOSONAR - value = None - else: - value = round(random.uniform(0, 100), 2) - - # Create data dictionary - data = { - 'tag': tag, - 'name': name, - 'timestamp': datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S'), - 'value': value - } - - # Send to Kafka - self.producer.send(topic, value=data) - - # Ensure all messages are sent - self.producer.flush() - - self.info("Success", metadata=metadata) diff --git a/scouter/activities/gates.py b/scouter/activities/gates.py index de8bf30..b5da7dd 100644 --- a/scouter/activities/gates.py +++ b/scouter/activities/gates.py @@ -1,20 +1,23 @@ -from temporalio import workflow, activity +from collections.abc import Hashable -from scouter import metrics +from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): - from sientia_do.notifications.models import NotificationLevel - from sientia_do.temporal.activities.base import BaseActivity - from sientia_do.notifications.handlers import NotificationHandler - from sientia_do.observability.logger import Logger - from scouter.utils.quality.filters import null_values_filter, out_of_bounds_filter - from typing import Any import traceback + from typing import Any + from pandas import DataFrame + from sientia_do.notifications.handlers import NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from sientia_do.observability.logger import Logger + from sientia_do.temporal.activities.base import BaseActivity + + from scouter import metrics + from scouter.utils.quality.filters import null_values_filter, out_of_bounds_filter quality_gate_filters = { 'NULL_VALUES_FILTER': null_values_filter, - 'OUT_OF_BOUNDS_FILTER': out_of_bounds_filter + 'OUT_OF_BOUNDS_FILTER': out_of_bounds_filter, } @@ -41,11 +44,11 @@ class Gates(BaseActivity): logger (Logger): Logger instance for operation logging notification_handler (NotificationHandler): Handler for system notifications """ - BaseActivity.__init__( - self, logger, notification_handler, set_error_counter=True) + BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True) - def apply_aggregation(self, group: DataFrame, aggr_function: str, - metadata: dict[str, Any]) -> float | None | str: + def apply_aggregation( + self, values: DataFrame, aggr_function: str, metadata: dict[str, Any] + ) -> float | None | str: """ Apply aggregation function to a group of time-series data. @@ -54,7 +57,7 @@ class Gates(BaseActivity): reporting for invalid aggregation functions. Args: - group (DataFrame): Group of data points to aggregate + values (DataFrame): Group of data points to aggregate (pre-sorted by timestamp) aggr_function (str): Aggregation function to apply. Supported functions: 'lts' (latest), 'avg' (average), 'mdn' (median), 'max' (maximum), 'min' (minimum) @@ -66,39 +69,42 @@ class Gates(BaseActivity): Raises: NotificationError: If invalid aggregation function is specified """ - if len(group) == 1: - return group['value'].item() + # Fast path for single value + if len(values) == 1: + return values['value'].iloc[0] - # Apply aggregation function to value if aggr_function == 'lts': - return group['value'].iloc[-1] + return values['value'].iloc[-1] + + # Remove NaN values without inplace operation + clean_values = values['value'].dropna() + + if clean_values.empty: + return None + + # Use dictionary lookup for aggregation functions (faster than if-elif chain) + aggregation_map = { + 'avg': lambda x: x.mean(), + 'mdn': lambda x: x.median(), + 'max': lambda x: x.max(), + 'min': lambda x: x.min(), + } + + if aggr_function in aggregation_map: + return aggregation_map[aggr_function](clean_values) else: - group.dropna(inplace=True, subset=['value']) + self.send_notification( + metadata=metadata, + notification_id='AGGREGATION_ISSUES', + message=f'Invalid aggregation function: {aggr_function}', + block='aggregate_data', + level=NotificationLevel.ERROR, + attachment_content=traceback.format_exc(), + ) + return 'continue' - if group.empty: - return None - - if aggr_function == 'avg': - return group['value'].mean() - elif aggr_function == 'mdn': - return group['value'].median() - elif aggr_function == 'max': - return group['value'].max() - elif aggr_function == 'min': - return group['value'].min() - else: - self.send_notification( - metadata=metadata, - notification_id="AGGREGATION_ISSUES", - message=f"Invalid aggregation function: {aggr_function}", - block="aggregate_data", - level=NotificationLevel.ERROR, - attachment_content=traceback.format_exc() - ) - return 'continue' - - @activity.defn(name="aggregate_data") - async def aggregate_data(self, input_data: dict[str, Any]) -> dict[str, Any]: + @activity.defn(name='aggregate_data') + async def aggregate_data(self, input_data: dict[str, Any]) -> dict[Hashable, Any]: """ Aggregate time-series data by tag and name using specified functions. @@ -113,7 +119,7 @@ class Gates(BaseActivity): - model_tags (dict[str, Any]): Tag configuration with aggregation functions Returns: - dict[str, Any]: Aggregated data organized by tag and name + dict[Hashable, Any]: Aggregated data organized by tag and name Raises: Exception: If aggregation operation fails @@ -125,93 +131,73 @@ class Gates(BaseActivity): # Convert input data to DataFrame df = DataFrame(input_data['data']) - self.info( - f"Aggregating time series data for {len(df)} rows", - metadata=metadata - ) + self.info(f'Aggregating time series data for {len(df)} rows', metadata=metadata) - # Initialize result dictionary - result = {} + # Sort once by timestamp for all data (more efficient than sorting each group) + df = df.sort_values(['tag', 'name', 'timestamp']) # Group by tag and name - grouped = df.groupby(['tag', 'name']) + # sort=False since we already sorted + grouped = df.groupby(['tag', 'name'], sort=False) + # Prepare aggregation functions mapping + model_tags = input_data['model_tags'] + + # Process groups efficiently + results = [] for (tag, name), group in grouped: # Get the aggregation function from model_tags - aggr_function = input_data['model_tags'].get( - name, {}).get('aggr_func', 'lts') + aggr_function = model_tags.get(name, {}).get('aggr_func', 'lts') - group.sort_values(by='timestamp', inplace=True) + # Get the latest timestamp (last row since data is sorted) + latest_timestamp = group['timestamp'].iloc[-1] - # Get the latest timestamp - latest_timestamp = group['timestamp'].max() - - aggr_value = self.apply_aggregation( - group, aggr_function, metadata) + aggr_value = self.apply_aggregation(group, aggr_function, metadata) if aggr_value == 'continue': continue - self.debug( - f"Aggregated data: {aggr_value}", - metadata=metadata - ) - self.debug( - f"Latest timestamp: {latest_timestamp}", - metadata=metadata - ) - self.debug( - f"Groups: {group.to_string()}", - metadata=metadata - ) - self.debug( - f"group name: {name}", - metadata=metadata - ) - self.debug( - f"group tag: {tag}", - metadata=metadata + # Store the result directly in list for better performance + results.append( + { + 'tag': tag, + 'name': name, + 'value': aggr_value, + 'timestamp': latest_timestamp, + 'aggregation_function': aggr_function, + } ) - # Store the result - result[f"{tag}_{name}"] = { - 'tag': tag, - 'name': name, - 'value': aggr_value, - 'timestamp': latest_timestamp, - 'aggregation_function': aggr_function - } + self.info(f'Aggregated data has {len(results)} rows', metadata=metadata) - result_df = DataFrame(list(result.values())) - self.info( - f"Aggregated data has {len(result_df)} rows", - metadata=metadata - ) + # Convert to DataFrame only once at the end if we have results + if results: + result_df = DataFrame(results) - self.debug( - f"Aggregated data: {result_df.to_string()}", - metadata=metadata - ) + self.debug(f'Final aggregated data:\n{result_df.to_string()}', metadata=metadata) - return result_df.to_dict() + return result_df.to_dict() + else: + # Return empty DataFrame dict structure + return DataFrame().to_dict() except Exception as e: trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="AGGREGATION_ISSUES", - message=f"Error aggregating data: {e}", - block="aggregate_data", + notification_id='AGGREGATION_ISSUES', + message=f'Error aggregating data: {e}', + block='aggregate_data', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) raise e - @activity.defn(name="data_quality_gate") - async def data_quality_gate(self, input_data: dict[str, Any]) -> dict[str, Any]: + @activity.defn(name='data_quality_gate') + async def data_quality_gate(self, input_data: dict[str, Any]) -> dict[Hashable, Any]: """ Apply data quality filters to incoming data. @@ -227,7 +213,7 @@ class Gates(BaseActivity): - model_tags (dict[str, Any]): Tag-specific validation rules Returns: - dict[str, Any]: Filtered data that passes quality validation + dict[Hashable, Any]: Filtered data that passes quality validation Raises: Exception: If quality validation fails @@ -239,10 +225,7 @@ class Gates(BaseActivity): data = DataFrame(input_data['data']) model_tags = input_data['model_tags'] - self.info( - f"Applying quality gate to data to {len(data)} rows", - metadata=metadata - ) + self.info(f'Applying quality gate to data to {len(data)} rows', metadata=metadata) tags = list(model_tags.keys()) @@ -251,25 +234,21 @@ class Gates(BaseActivity): for filter_name, config in filters.items(): policy = config['policy'] if filter_name not in quality_gate_filters: - self.warning( - f"Filter {filter_name} not found", - metadata=metadata - ) + self.warning(f'Filter {filter_name} not found', metadata=metadata) continue try: - filtered_data = quality_gate_filters[filter_name]( - data, model_tags) + filtered_data = quality_gate_filters[filter_name](data, model_tags) except Exception as e: trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="DATA_QUALITY_GATE_ISSUES", - message=f"Error applying filter {filter_name}: {e}", - block="data_quality_gate", + notification_id='DATA_QUALITY_GATE_ISSUES', + message=f'Error applying filter {filter_name}: {e}', + block='data_quality_gate', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) self.error(trace, metadata=metadata) @@ -278,30 +257,27 @@ class Gates(BaseActivity): if filtered_data.empty: continue - message = f"{len(filtered_data)} rows has quality issues: {filter_name}: {policy}" + message = f'{len(filtered_data)} rows has quality issues: {filter_name}: {policy}' attachment = filtered_data.to_string() self.send_notification( metadata=metadata, - notification_id=f"DATA_QUALITY_GATE_ISSUES__{filter_name}", + notification_id=f'DATA_QUALITY_GATE_ISSUES__{filter_name}', message=message, - block="data_quality_gate", + block='data_quality_gate', level=NotificationLevel.WARNING, - attachment_content=attachment + attachment_content=attachment, ) - if policy == "DISCARD": + if policy == 'DISCARD': data = data[~data.index.isin(filtered_data.index)] - self.info( - f"Data quality gate applied, final data has {len(data)} rows", - metadata=metadata - ) + self.info(f'Data quality gate applied, final data has {len(data)} rows', metadata=metadata) return data.to_dict() - @activity.defn(name="write_metrics") - async def write_metrics(self, input_data: dict[str, Any]): + @activity.defn(name='write_metrics') + async def write_metrics(self, input_data: dict[str, Any]) -> None: """ Write metrics to the database. input_data: @@ -309,18 +285,12 @@ class Gates(BaseActivity): """ metadata = input_data['metadata'] - self.info( - f"Writing metrics for {metadata['model_name']}", - metadata=metadata - ) + self.info(f'Writing metrics for {metadata["model_name"]}', metadata=metadata) metrics.LABORIOUS_DATA_WRITTEN_COUNT.labels( pod_id=self.pod_id, model_name=metadata['model_name'], - pipeline_name=metadata['workflow_name'] + pipeline_name=metadata['workflow_name'], ).inc() - self.info( - f"Metrics written for {metadata['model_name']}", - metadata=metadata - ) + self.info(f'Metrics written for {metadata["model_name"]}', metadata=metadata) diff --git a/scouter/activities/mongodb.py b/scouter/activities/mongodb.py index 3c5b9f4..a20e356 100644 --- a/scouter/activities/mongodb.py +++ b/scouter/activities/mongodb.py @@ -1,15 +1,19 @@ -from temporalio import workflow, activity +from collections.abc import Hashable +from datetime import UTC + +from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): - from typing import Any import traceback - from datetime import datetime, timezone - from pymongo import MongoClient + from datetime import datetime + from typing import Any + from pandas import DataFrame + from pymongo import MongoClient from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.models import NotificationLevel - from sientia_do.temporal.activities.base import BaseActivity from sientia_do.observability.logger import Logger + from sientia_do.temporal.activities.base import BaseActivity from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ @@ -35,10 +39,10 @@ def clear_mongo_id(docs: list) -> list: clear_mongo_id(doc) elif isinstance(doc, dict): - if "_id" in doc: - del doc["_id"] + if '_id' in doc: + del doc['_id'] - for key, value in doc.items(): + for _key, value in doc.items(): if isinstance(value, list): clear_mongo_id(value) elif isinstance(value, dict): @@ -62,9 +66,13 @@ class MongoDB(BaseActivity): distributed data processing with fault tolerance and monitoring. """ - def __init__(self, connection_string: str, database_name: str, - logger: Logger, - notification_handler: NotificationHandler): + def __init__( + self, + connection_string: str, + database_name: str, + logger: Logger, + notification_handler: NotificationHandler, + ): """ Initialize MongoDB connection and services. @@ -80,19 +88,19 @@ class MongoDB(BaseActivity): self.connection_string = connection_string self.database_name = database_name - self.client = MongoClient( - self.connection_string, serverSelectionTimeoutMS=5000) + self.client: MongoClient = MongoClient( + self.connection_string, serverSelectionTimeoutMS=5000 + ) self.client.server_info() # Trigger an exception if connection fails self.database = self.client[self.database_name] # Initialize MongoDB client here (omitted for brevity) - logger.info("MongoDB connection initialized") + logger.info('MongoDB connection initialized') - BaseActivity.__init__(self, - logger=logger, - notification_handler=notification_handler, - set_error_counter=True) + BaseActivity.__init__( + self, logger=logger, notification_handler=notification_handler, set_error_counter=True + ) def shutdown(self): """ @@ -103,11 +111,11 @@ class MongoDB(BaseActivity): """ try: if self.client: - self.logger.info("Closing MongoDB connection...") + self.logger.info('Closing MongoDB connection...') self.client.close() - self.logger.info("MongoDB connection closed successfully") + self.logger.info('MongoDB connection closed successfully') except Exception as e: - self.logger.error(f"Failed to close MongoDB connection: {e}") + self.logger.error(f'Failed to close MongoDB connection: {e}') def __del__(self): """ @@ -118,8 +126,8 @@ class MongoDB(BaseActivity): """ self.shutdown() - @activity.defn(name="load_latest_data") - async def load_latest_data(self, input_data: dict[str, Any]) -> dict[str, Any]: + @activity.defn(name='load_latest_data') + async def load_latest_data(self, input_data: dict[str, Any]) -> dict[Hashable, Any]: """ Load the latest data from MongoDB collection since a specified timestamp. @@ -144,60 +152,44 @@ class MongoDB(BaseActivity): collection_name = input_data['collection_name'] last_data_timestamp = input_data['last_data_timestamp'] - self.info( - f"Loading data from MongoDB: {input_data}", - metadata=metadata - ) + self.info(f'Loading data from MongoDB: {input_data}', metadata=metadata) try: - if last_data_timestamp is None: data_filter = {} else: data_filter = { - "inserted_at": { - "$gt": datetime.strptime(last_data_timestamp, DATETIME_FORMAT_MS_WITH_TZ) + 'inserted_at': { + '$gt': datetime.strptime(last_data_timestamp, DATETIME_FORMAT_MS_WITH_TZ) } } - self.debug( - f"Data filter: {data_filter}", - metadata=metadata - ) + self.debug(f'Data filter: {data_filter}', metadata=metadata) - data = list(self.database[collection_name].find( - data_filter, {"_id": 0})) + data = list(self.database[collection_name].find(data_filter, {'_id': 0})) data = clear_mongo_id(data) - self.debug( - f"Collected: {data}", - metadata=metadata - ) + self.debug(f'Collected: {data}', metadata=metadata) for item in data: - item['inserted_at'] = item['inserted_at'].replace( - tzinfo=timezone.utc).strftime(DATETIME_FORMAT_MS_WITH_TZ) + item['inserted_at'] = ( + item['inserted_at'].replace(tzinfo=UTC).strftime(DATETIME_FORMAT_MS_WITH_TZ) + ) - self.info( - f"Loaded {len(data)} documents from MongoDB", - metadata=metadata - ) + self.info(f'Loaded {len(data)} documents from MongoDB', metadata=metadata) - self.debug( - f"Loaded data: {data}", - metadata=metadata - ) + self.debug(f'Loaded data: {data}', metadata=metadata) return DataFrame(data).to_dict() except Exception as e: trace = traceback.format_exc() self.send_notification( metadata=metadata, - notification_id="MONGO_LOAD_ERROR", - message=f"Error loading data from MongoDB: {e}", - block="load_latest_data", + notification_id='MONGO_LOAD_ERROR', + message=f'Error loading data from MongoDB: {e}', + block='load_latest_data', level=NotificationLevel.ERROR, - attachment_content=trace + attachment_content=trace, ) raise e diff --git a/scouter/activities/redis.py b/scouter/activities/redis.py index c8a4877..d724772 100644 --- a/scouter/activities/redis.py +++ b/scouter/activities/redis.py @@ -1,17 +1,20 @@ -from temporalio import workflow, activity +from collections.abc import Hashable + +from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): - from logging import Logger import traceback + from typing import Any + + from pandas import DataFrame from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.models import NotificationLevel - from sientia_do.temporal.activities.redis_base import Redis as RedisBase from sientia_do.observability.logger import Logger - from typing import Any - from pandas import DataFrame - from scouter import metrics + from sientia_do.temporal.activities.redis_base import Redis as RedisBase from sientia_do.temporal.constants import DATETIME_FORMAT, now + from scouter import metrics + class Redis(RedisBase): """ @@ -28,9 +31,15 @@ class Redis(RedisBase): distributed data processing with fault tolerance and monitoring. """ - def __init__(self, host: str, port: int, - username: str, password: str, - logger: Logger, notification_handler: NotificationHandler): + def __init__( + self, + host: str, + port: int, + username: str, + password: str, + logger: Logger, + notification_handler: NotificationHandler, + ): """ Initialize Redis connection and services. @@ -42,10 +51,9 @@ class Redis(RedisBase): logger (Logger): Logger instance for operation logging notification_handler (NotificationHandler): Handler for system notifications """ - RedisBase.__init__(self, host, port, username, - password, logger, notification_handler) + RedisBase.__init__(self, host, port, username, password, logger, notification_handler) - @activity.defn(name="get_last_data_timestamp") + @activity.defn(name='get_last_data_timestamp') async def get_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None: """ Retrieve the last processed data timestamp from Redis. @@ -68,34 +76,31 @@ class Redis(RedisBase): Exception: If Redis operation fails """ metadata = input_data['metadata'] - key = f"last_data_timestamp:{input_data['workflow_name']}:{input_data['schedule_name']}" + key = f'last_data_timestamp:{input_data["workflow_name"]}:{input_data["schedule_name"]}' - self.info(f"Getting last data timestamp for {key}") + self.info(f'Getting last data timestamp for {key}') try: data_hold = self.get(key) except Exception as e: self.send_notification( metadata=metadata, - notification_id="REDIS_GET_ERROR", - message=f"Error getting last data timestamp: {e}", - block="get_last_data_timestamp", + notification_id='REDIS_GET_ERROR', + message=f'Error getting last data timestamp: {e}', + block='get_last_data_timestamp', level=NotificationLevel.ERROR, - attachment_content=traceback.format_exc() + attachment_content=traceback.format_exc(), ) raise e - self.info( - f"Last collected timestamp: {data_hold}", - metadata=metadata - ) + self.info(f'Last collected timestamp: {data_hold}', metadata=metadata) if not data_hold: return None return data_hold - @activity.defn(name="put_last_data_timestamp") + @activity.defn(name='put_last_data_timestamp') async def put_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None: """ Store the last processed data timestamp in Redis. @@ -119,43 +124,37 @@ class Redis(RedisBase): Exception: If Redis operation fails """ metadata = input_data['metadata'] - key = f"last_data_timestamp:{input_data['workflow_name']}:{input_data['schedule_name']}" + key = f'last_data_timestamp:{input_data["workflow_name"]}:{input_data["schedule_name"]}' - self.info(f"Putting last data timestamp for {key}") + self.info(f'Putting last data timestamp for {key}') data = DataFrame(input_data['data']) if data.empty: - self.warning("No data to insert", - metadata=metadata - ) + self.warning('No data to insert', metadata=metadata) return None last_data_timestamp = data['inserted_at'].max() - self.info( - f"Last collected timestamp to insert: {last_data_timestamp}", - metadata=metadata - ) + self.info(f'Last collected timestamp to insert: {last_data_timestamp}', metadata=metadata) try: - self.set(key, last_data_timestamp, ttl=60*60*5) + self.set(key, last_data_timestamp, ttl=60 * 60 * 5) except Exception as e: self.send_notification( metadata=metadata, - notification_id="REDIS_SET_ERROR", - message=f"Error setting last data timestamp: {e}", - - block="put_last_data_timestamp", + notification_id='REDIS_SET_ERROR', + message=f'Error setting last data timestamp: {e}', + block='put_last_data_timestamp', level=NotificationLevel.ERROR, - attachment_content=traceback.format_exc() + attachment_content=traceback.format_exc(), ) raise e return last_data_timestamp - @activity.defn(name="group_and_hold_data") - async def group_and_hold_data(self, input_data: dict[str, Any]) -> dict[str, Any]: + @activity.defn(name='group_and_hold_data') + async def group_and_hold_data(self, input_data: dict[str, Any]) -> dict[Hashable, Any]: """ Group data by tags and store temporarily in Redis with TTL. @@ -182,109 +181,92 @@ class Redis(RedisBase): """ metadata = input_data['metadata'] - self.debug("Grouping and holding data...", - metadata=metadata - ) + self.debug('Grouping and holding data...', metadata=metadata) data = DataFrame(input_data['data']) model_tags = input_data['model_tags'] retention_time = input_data['retention_time'] - key = f"held_data_{input_data['workflow_name']}_{input_data['schedule_name']}" + key = f'held_data_{input_data["workflow_name"]}_{input_data["schedule_name"]}' - self.info(f"Getting held data for {key}") + self.info(f'Getting held data for {key}') try: data_hold = self.get(key) except Exception as e: self.send_notification( metadata=metadata, - notification_id="REDIS_GET_ERROR", - message=f"Error getting held data: {e}", - block="group_and_hold_data", + notification_id='REDIS_GET_ERROR', + message=f'Error getting held data: {e}', + block='group_and_hold_data', level=NotificationLevel.ERROR, - attachment_content=traceback.format_exc() + attachment_content=traceback.format_exc(), ) raise e if not data_hold: data_hold = {} if data.empty: - self.warning("No data to export", - metadata=metadata - ) + self.warning('No data to export', metadata=metadata) return data_hold - self.info(f"Grouping and holding data for {len(data)} rows") + self.info(f'Grouping and holding data for {len(data)} rows') try: - # Remove possibly removed tags tags = list(model_tags.keys()) tags.append('timestamp') - self.debug( - f"Tags to keep: {tags}", - metadata=metadata - ) - data_hold = {tag: content for tag, - content in data_hold.items() if tag in tags} - self.debug( - f"Data hold after removing removed tags: {data_hold}", - metadata=metadata - ) + self.debug(f'Tags to keep: {tags}', metadata=metadata) + data_hold = {tag: content for tag, content in data_hold.items() if tag in tags} + self.debug(f'Data hold after removing removed tags: {data_hold}', metadata=metadata) to_register_metrics = [] for _, row in data.iterrows(): value = row['value'] data_hold[row['name']] = value - to_register_metrics.append( - (row['name'], value)) + to_register_metrics.append((row['name'], value)) - data_hold['timestamp'] = data['timestamp'].max() if not data.empty else \ - data_hold['timestamp'] + data_hold['timestamp'] = ( + data['timestamp'].max() if not data.empty else data_hold['timestamp'] + ) self.set(key, data_hold, ttl=retention_time) # Register metrics - self.debug( - f"Metrics to register: {to_register_metrics}", - metadata=metadata - ) + self.debug(f'Metrics to register: {to_register_metrics}', metadata=metadata) for metric in to_register_metrics: metrics.TAG_CHANGES_MONITOR.labels( pod_id=self.pod_id, model_name=metadata['model_name'], pipeline_name=metadata['workflow_name'], - tag_name=metric[0] + tag_name=metric[0], ).set(metric[1]) data_hold_df = DataFrame(data_hold, index=[0]) data_hold_melted = data_hold_df.melt( - id_vars='timestamp', var_name='variable', value_name='value') + id_vars='timestamp', var_name='variable', value_name='value' + ) data_hold_melted['model_id'] = input_data['model_id'] data_hold_melted.reset_index(drop=True, inplace=True) except Exception as e: self.send_notification( metadata=metadata, - notification_id="REDIS_SET_ERROR", - message=f"Error setting held data: {e}", - block="group_and_hold_data", + notification_id='REDIS_SET_ERROR', + message=f'Error setting held data: {e}', + block='group_and_hold_data', level=NotificationLevel.ERROR, - attachment_content=traceback.format_exc() + attachment_content=traceback.format_exc(), ) raise e - self.info(f"Data held and melted has {len(data_hold_melted)} rows") + self.info(f'Data held and melted has {len(data_hold_melted)} rows') - self.debug( - f"Data held and melted:\n {data_hold_melted.to_string()}", - metadata=metadata - ) + self.debug(f'Data held and melted:\n {data_hold_melted.to_string()}', metadata=metadata) return data_hold_melted.to_dict() - @activity.defn(name="store_data_package") + @activity.defn(name='store_data_package') async def store_data_package(self, input_data: dict[str, Any]): """ Stores the data package in redis. It's a debug feature and must be toggled on. @@ -296,25 +278,22 @@ class Redis(RedisBase): data: The data used to collect the data. """ metadata = input_data['metadata'] - key = f"data_package_{input_data['workflow_name']}_{input_data['schedule_name']}_{now().strftime(DATETIME_FORMAT)}" + key = f'data_package_{input_data["workflow_name"]}_{input_data["schedule_name"]}_{now().strftime(DATETIME_FORMAT)}' data = DataFrame(input_data['data']) held_data = DataFrame(input_data['held_data']) - cache = { - 'data': data.to_dict(), - 'held_data': held_data.to_dict() - } + cache = {'data': data.to_dict(), 'held_data': held_data.to_dict()} try: self.set(key, cache, ttl=120) except Exception as e: self.send_notification( metadata=metadata, - notification_id="REDIS_SET_ERROR", - message=f"Error setting data package: {e}", - block="store_data_package", + notification_id='REDIS_SET_ERROR', + message=f'Error setting data package: {e}', + block='store_data_package', level=NotificationLevel.ERROR, - attachment_content=traceback.format_exc() + attachment_content=traceback.format_exc(), ) raise e diff --git a/scouter/metrics.py b/scouter/metrics.py index a1d3e07..ac249c7 100644 --- a/scouter/metrics.py +++ b/scouter/metrics.py @@ -1,25 +1,25 @@ -from prometheus_client import Gauge, Counter +from prometheus_client import Counter, Gauge # Application health and status metrics APP_UP = Gauge( - "app_up", - "Indicates if the application is running (1) or shutting down (0)", - ["pod_id"], + 'app_up', + 'Indicates if the application is running (1) or shutting down (0)', + ['pod_id'], ) # Core labels for consistent metric labeling -CORE_LABELS = ["pod_id", "model_name", "pipeline_name"] +CORE_LABELS = ['pod_id', 'model_name', 'pipeline_name'] # Data processing metrics LABORIOUS_DATA_WRITTEN_COUNT = Counter( - "scouter_laborious_data_written_count", - "Number of writings to the database table laborious_data", + 'scouter_laborious_data_written_count', + 'Number of writings to the database table laborious_data', CORE_LABELS, ) # Tag monitoring metrics TAG_CHANGES_MONITOR = Gauge( - "scouter_tag_changes_monitor", - "Current value change of each tag", - [*CORE_LABELS, "tag_name"], + 'scouter_tag_changes_monitor', + 'Current value change of each tag', + [*CORE_LABELS, 'tag_name'], ) diff --git a/scouter/utils/connectors_config.py b/scouter/utils/connectors_config.py index 176fce4..4264aa1 100644 --- a/scouter/utils/connectors_config.py +++ b/scouter/utils/connectors_config.py @@ -23,7 +23,7 @@ def build_postgres_config() -> dict[str, Any]: 'password': getenv('POSTGRES_PASSWORD', 'sientia'), 'dbname': getenv('POSTGRES_DBNAME', 'sientia'), 'min_connections': int(getenv('POSTGRES_MIN_CONNECTIONS', '5')), - 'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20')) + 'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20')), } @@ -40,7 +40,7 @@ def build_kafka_config() -> dict[str, Any]: return { 'bootstrap_servers': getenv('KAFKA_BOOTSTRAP_SERVERS', 'localhost:9092'), 'polling_time': int(getenv('KAFKA_POLLING_TIME', '1000')), - 'group_id': 'scouter-group' + 'group_id': 'scouter-group', } @@ -59,7 +59,7 @@ def build_redis_config() -> dict[str, Any]: 'host': getenv('REDIS_HOST', 'localhost'), 'port': int(getenv('REDIS_PORT', '6379')), 'username': getenv('REDIS_USERNAME', None), - 'password': getenv('REDIS_PASSWORD', None) + 'password': getenv('REDIS_PASSWORD', None), } @@ -79,7 +79,7 @@ def build_mongodb_config() -> dict[str, Any]: connection_string = f'mongodb://{username}:{password}@{uri}' return { 'connection_string': connection_string, - 'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia') + 'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'), } diff --git a/scouter/utils/quality/filters.py b/scouter/utils/quality/filters.py index 4586b2a..5c466d7 100644 --- a/scouter/utils/quality/filters.py +++ b/scouter/utils/quality/filters.py @@ -1,7 +1,8 @@ -from pandas import DataFrame -import numpy as np from typing import Any +import numpy as np +from pandas import DataFrame + def check_data_range(value: float | int | None, val_range: list) -> bool: """ @@ -49,9 +50,14 @@ def out_of_bounds_filter(df: DataFrame, model_tags: dict[str, Any]) -> DataFrame Note: Tags without data_range configuration are treated as having infinite bounds """ - return df[df.apply(lambda x: check_data_range( - x['value'], model_tags[x['name']].get('data_range', (-np.inf, np.inf))), - axis=1)] + return df[ + df.apply( + lambda x: check_data_range( + x['value'], model_tags[x['name']].get('data_range', (-np.inf, np.inf)) + ), + axis=1, + ) + ] def null_values_filter(df: DataFrame, _model_tags: dict[str, Any]) -> DataFrame: diff --git a/scouter/worker/worker.py b/scouter/worker/worker.py index c0d1066..c842761 100644 --- a/scouter/worker/worker.py +++ b/scouter/worker/worker.py @@ -1,29 +1,29 @@ -from temporalio import workflow, client -from temporalio.worker import Worker, PollerBehaviorAutoscaling -from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig +from temporalio import client, workflow +from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig +from temporalio.worker import PollerBehaviorAutoscaling, Worker with workflow.unsafe.imports_passed_through(): - import sys + import asyncio import os + import sys + + from prometheus_client import start_http_server from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.observability.logger import get_logger - from scouter.activities.activities import Activities - from scouter.workflow.scouter import Scouter - from scouter.workflow.sub_workflows.core_scouter import CoreScouter - from scouter.workflow.fake_data import FakeData - from scouter.activities.faker import Faker - import asyncio - from prometheus_client import start_http_server + from scouter import metrics + from scouter.activities.activities import Activities from scouter.utils.connectors_config import ( + build_mongodb_config, build_postgres_config, build_redis_config, - build_mongodb_config ) + from scouter.workflow.scouter import Scouter + from scouter.workflow.sub_workflows.core_scouter import CoreScouter # Environment configuration -POD_ID = os.getenv("HOSTNAME", "localhost") -SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091")) +POD_ID = os.getenv('HOSTNAME', 'localhost') +SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091')) async def main(): @@ -59,9 +59,9 @@ async def main(): 'schedule_name': '-', } - logger.custom_info(f"Starting Worker with pod_id: {POD_ID}", metadata) + logger.custom_info(f'Starting Worker with pod_id: {POD_ID}', metadata) - logger.custom_info("Starting prometheus client...", metadata) + logger.custom_info('Starting prometheus client...', metadata) start_prometheus_server() logger.custom_info('Starting Notification Handler...', metadata) @@ -71,7 +71,7 @@ async def main(): connection_string=mongo_config['connection_string'], database=mongo_config['database_name'], logger=logger, - project_name=os.getenv('PROJECT_NAME', 'scouter') + project_name=os.getenv('PROJECT_NAME', 'scouter'), ) logger.custom_info('Starting Activities...', metadata) @@ -81,34 +81,21 @@ async def main(): notification_handler=notification_handler, postgres_config=build_postgres_config(), redis_config=build_redis_config(), - mongodb_config=build_mongodb_config() + mongodb_config=build_mongodb_config(), ) - logger.custom_info('Starting Faker Activities...', metadata) - - faker_activities = Faker( - logger=logger, - notification_handler=notification_handler, - bootstrap_servers=os.getenv( - 'KAFKA_BOOTSTRAP_SERVERS', 'localhost:9092') - ) - - logger.custom_info( - f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata) + logger.custom_info(f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata) new_runtime = Runtime( telemetry=TelemetryConfig( - metrics=PrometheusConfig( - bind_address=f"0.0.0.0:{SDK_METRICS_PORT}") + metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}') ) ) logger.custom_info('Starting Temporal Client...', metadata) temporal_client = await client.Client.connect( - target_host=host, - namespace=os.getenv('TEMPORAL_NAMESPACE', 'scouter'), - runtime=new_runtime + target_host=host, namespace=os.getenv('TEMPORAL_NAMESPACE', 'scouter'), runtime=new_runtime ) logger.custom_info('Starting Workers...', metadata) @@ -134,15 +121,7 @@ async def main(): max_concurrent_local_activities=50, max_cached_workflows=200, workflow_task_poller_behavior=PollerBehaviorAutoscaling(), - activity_task_poller_behavior=PollerBehaviorAutoscaling() - ), - Worker( - temporal_client, - task_queue='fake_data-queue', - workflows=[FakeData], - activities=[ - faker_activities.generate_and_send_data, - ] + activity_task_poller_behavior=PollerBehaviorAutoscaling(), ) ] @@ -156,8 +135,9 @@ async def main(): await asyncio.gather(*handlers) except BaseException as e: # NOSONAR - logger.custom_error("An unhandled exception occurred: %s", - e, exc_info=True, metadata=metadata) + logger.custom_error( + 'An unhandled exception occurred: %s', e, exc_info=True, metadata=metadata + ) finally: if notification_handler: notification_handler.shutdown() @@ -183,12 +163,12 @@ def start_prometheus_server(): SystemExit: If metrics server fails to start """ try: - port = int(os.getenv("HTTP_METRICS_PORT", 9090)) + port = int(os.getenv('HTTP_METRICS_PORT', 9090)) start_http_server(port) - print(f"Prometheus server started on port {port}.") + print(f'Prometheus server started on port {port}.') metrics.APP_UP.labels(pod_id=POD_ID).set(1) # Mark app as UP except Exception as e: - print(f"Failed to start Prometheus server: {e}") + print(f'Failed to start Prometheus server: {e}') os._exit(1) diff --git a/scouter/workflow/fake_data.py b/scouter/workflow/fake_data.py deleted file mode 100644 index 95cdde1..0000000 --- a/scouter/workflow/fake_data.py +++ /dev/null @@ -1,57 +0,0 @@ -from temporalio import workflow - -with workflow.unsafe.imports_passed_through(): - from scouter.activities.faker import Faker - from datetime import timedelta - from typing import Dict, Any - from sientia_do.temporal.policies import retry_policy - - -@workflow.defn(name="fake_data") -class FakeData: - """ - Test data generation workflow for development and testing purposes. - - This workflow generates synthetic industrial sensor data and publishes it to - Kafka topics. It's designed for: - - Development and testing of data processing pipelines - - Load testing of downstream systems - - Demonstration of data flow patterns - - Validation of data quality filters and aggregation functions - - The generated data simulates realistic industrial sensor readings with - configurable message counts and topic routing. - """ - - @workflow.run - async def run(self, workflow_input: Dict[str, Any]) -> str: - """ - Execute the fake data generation workflow. - - This method generates synthetic sensor data and publishes it to the specified - Kafka topic. The data includes realistic industrial sensor readings with - configurable parameters for testing and development purposes. - - Args: - workflow_input (dict[str, Any]): Workflow configuration parameters. - Required fields: - - topic (str): Kafka topic name for data publication - - metadata (dict[str, Any], optional): Workflow execution metadata - - num_messages (int, optional): Number of messages to generate. - Defaults to random count between 1 and available sensor tags. - - Returns: - str: Success confirmation message - - Raises: - WorkflowExecutionError: If workflow execution fails - ActivityExecutionError: If data generation or Kafka publishing fails - """ - await workflow.execute_activity_method( - Faker.generate_and_send_data, - { - 'topic': workflow_input['topic'] - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) - ) diff --git a/scouter/workflow/scouter.py b/scouter/workflow/scouter.py index 0acb095..e028bc6 100644 --- a/scouter/workflow/scouter.py +++ b/scouter/workflow/scouter.py @@ -1,13 +1,15 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from scouter.activities.activities import Activities - from typing import Any from datetime import timedelta + from typing import Any + from sientia_do.temporal.policies import retry_policy + from scouter.activities.activities import Activities -@workflow.defn(name="scouter") + +@workflow.defn(name='scouter') class Scouter: """ Main Scouter workflow that orchestrates data ingestion and processing. @@ -67,7 +69,7 @@ class Scouter: 'model_id': input_data['model_id'], 'model_name': input_data['model_name'], 'schedule_name': input_data['schedule_name'], - 'workflow_name': input_data['workflow_name'] + 'workflow_name': input_data['workflow_name'], } } @@ -76,21 +78,21 @@ class Scouter: { **metadata, 'workflow_name': input_data['workflow_name'], - 'schedule_name': input_data['schedule_name'] + 'schedule_name': input_data['schedule_name'], }, start_to_close_timeout=timedelta(seconds=60), - retry_policy=retry_policy + retry_policy=retry_policy, ) data = await workflow.execute_local_activity_method( Activities.load_latest_data, { **metadata, - 'collection_name': f"raw_{input_data['schedule_name']}", - 'last_data_timestamp': last_data_timestamp + 'collection_name': f'raw_{input_data["schedule_name"]}', + 'last_data_timestamp': last_data_timestamp, }, start_to_close_timeout=timedelta(seconds=60), - retry_policy=retry_policy + retry_policy=retry_policy, ) if data == {}: @@ -102,16 +104,13 @@ class Scouter: **metadata, 'data': data, 'workflow_name': input_data['workflow_name'], - 'schedule_name': input_data['schedule_name'] + 'schedule_name': input_data['schedule_name'], }, start_to_close_timeout=timedelta(seconds=60), - retry_policy=retry_policy + retry_policy=retry_policy, ) input_data['data'] = data input_data['metadata'] = metadata - await workflow.execute_child_workflow( - 'core_scouter', - input_data - ) + await workflow.execute_child_workflow('core_scouter', input_data) diff --git a/scouter/workflow/sub_workflows/core_scouter.py b/scouter/workflow/sub_workflows/core_scouter.py index f4586c7..1531fb9 100644 --- a/scouter/workflow/sub_workflows/core_scouter.py +++ b/scouter/workflow/sub_workflows/core_scouter.py @@ -1,14 +1,16 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from scouter.activities.activities import Activities - from typing import Any from datetime import timedelta - from sientia_do.temporal.policies import retry_policy + from typing import Any + from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + from sientia_do.temporal.policies import retry_policy + + from scouter.activities.activities import Activities -@workflow.defn(name="core_scouter") +@workflow.defn(name='core_scouter') class CoreScouter: """ Core data processing workflow that handles data quality, aggregation, and export. @@ -68,21 +70,17 @@ class CoreScouter: **metadata, 'filters': input_data['filters'], 'data': input_data['data'], - 'model_tags': input_data['model_tags'] + 'model_tags': input_data['model_tags'], }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) grouped_data = await workflow.execute_local_activity_method( Activities.aggregate_data, - { - **metadata, - 'data': filtered_data, - 'model_tags': input_data['model_tags'] - }, + {**metadata, 'data': filtered_data, 'model_tags': input_data['model_tags']}, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) held_data = await workflow.execute_local_activity_method( @@ -94,10 +92,10 @@ class CoreScouter: 'data': grouped_data, 'model_id': input_data['model_id'], 'model_tags': input_data['model_tags'], - 'retention_time': input_data['retention_time'] + 'retention_time': input_data['retention_time'], }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) if held_data == {}: @@ -110,13 +108,10 @@ class CoreScouter: 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': held_data, - 'timestamp_conversion': { - 'column': 'timestamp', - 'format': DATETIME_FORMAT_WITH_TZ - } + 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) await workflow.execute_activity_method( @@ -125,7 +120,7 @@ class CoreScouter: **metadata, }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) if input_data.get('debug_data_package', False): @@ -139,5 +134,5 @@ class CoreScouter: 'schedule_name': input_data['schedule_name'], }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) + start_to_close_timeout=timedelta(seconds=60), ) diff --git a/tests/activities/test_activities.py b/tests/activities/test_activities.py index ecb5c8d..2ec9bb5 100644 --- a/tests/activities/test_activities.py +++ b/tests/activities/test_activities.py @@ -1,10 +1,11 @@ -from unittest.mock import patch, MagicMock, ANY -from pytest import mark +from unittest.mock import ANY, MagicMock, patch + from sientia_do.temporal.activities.postgres import Postgres + from scouter.activities.activities import Activities +from scouter.activities.gates import Gates from scouter.activities.mongodb import MongoDB from scouter.activities.redis import Redis -from scouter.activities.gates import Gates @patch('scouter.activities.activities.MongoDB.__init__') @@ -12,7 +13,6 @@ from scouter.activities.gates import Gates @patch('scouter.activities.activities.Redis.__init__') @patch('scouter.activities.activities.Gates.__init__') def test___init__(mock_gates_init, mock_redis_init, mock_postgres_init, mock_mongodb_init): - postgres_config = { 'host': 'localhost', 'port': 5432, @@ -20,19 +20,14 @@ def test___init__(mock_gates_init, mock_redis_init, mock_postgres_init, mock_mon 'password': 'postgres', 'dbname': 'postgres', 'min_connections': 1, - 'max_connections': 10 + 'max_connections': 10, } - redis_config = { - 'host': 'localhost', - 'port': 6379, - 'username': 'redis', - 'password': 'redis' - } + redis_config = {'host': 'localhost', 'port': 6379, 'username': 'redis', 'password': 'redis'} mongodb_config = { 'connection_string': 'mongodb://localhost:27017', - 'database_name': 'test_database' + 'database_name': 'test_database', } logger = MagicMock() @@ -43,7 +38,7 @@ def test___init__(mock_gates_init, mock_redis_init, mock_postgres_init, mock_mon redis_config=redis_config, mongodb_config=mongodb_config, logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) assert isinstance(activities, Activities) @@ -62,7 +57,7 @@ def test___init__(mock_gates_init, mock_redis_init, mock_postgres_init, mock_mon min_connections=postgres_config['min_connections'], max_connections=postgres_config['max_connections'], logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) mock_redis_init.assert_called_once_with( @@ -72,7 +67,7 @@ def test___init__(mock_gates_init, mock_redis_init, mock_postgres_init, mock_mon username=redis_config['username'], password=redis_config['password'], logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) mock_mongodb_init.assert_called_once_with( @@ -80,13 +75,11 @@ def test___init__(mock_gates_init, mock_redis_init, mock_postgres_init, mock_mon connection_string=mongodb_config['connection_string'], database_name=mongodb_config['database_name'], logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) mock_gates_init.assert_called_once_with( - ANY, - logger=logger, - notification_handler=notification_handler + ANY, logger=logger, notification_handler=notification_handler ) @@ -96,8 +89,14 @@ def test___init__(mock_gates_init, mock_redis_init, mock_postgres_init, mock_mon @patch('scouter.activities.activities.MongoDB.__init__') @patch('scouter.activities.activities.Postgres.close') @patch('scouter.activities.activities.MongoDB.shutdown') -def test_shutdown(mock_mongodb_close, mock_postgres_close, _mock_mongodb_init, - _mock_gates_init, _mock_redis_init, _mock_postgres_init): +def test_shutdown( + mock_mongodb_close, + mock_postgres_close, + _mock_mongodb_init, + _mock_gates_init, + _mock_redis_init, + _mock_postgres_init, +): postgres_config = { 'host': 'localhost', 'port': 5432, @@ -105,19 +104,14 @@ def test_shutdown(mock_mongodb_close, mock_postgres_close, _mock_mongodb_init, 'password': 'postgres', 'dbname': 'postgres', 'min_connections': 1, - 'max_connections': 10 + 'max_connections': 10, } - redis_config = { - 'host': 'localhost', - 'port': 6379, - 'username': 'redis', - 'password': 'redis' - } + redis_config = {'host': 'localhost', 'port': 6379, 'username': 'redis', 'password': 'redis'} mongodb_config = { 'connection_string': 'mongodb://localhost:27017', - 'database_name': 'test_database' + 'database_name': 'test_database', } logger = MagicMock() @@ -128,7 +122,7 @@ def test_shutdown(mock_mongodb_close, mock_postgres_close, _mock_mongodb_init, redis_config=redis_config, mongodb_config=mongodb_config, logger=logger, - notification_handler=notification_handler + notification_handler=notification_handler, ) activities.shutdown() diff --git a/tests/activities/test_faker.py b/tests/activities/test_faker.py deleted file mode 100644 index 0e3a61c..0000000 --- a/tests/activities/test_faker.py +++ /dev/null @@ -1,153 +0,0 @@ -from unittest.mock import MagicMock, patch, call -import pytest -from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler -from scouter.activities.faker import Faker - - -@pytest.fixture -def mock_kafka_producer(): - with patch('scouter.activities.faker.KafkaProducer') as mock: - producer = MagicMock() - mock.return_value = producer - yield producer - - -@pytest.fixture -def mock_datetime(): - with patch('scouter.activities.faker.datetime') as mock_dt: - mock_dt.now.return_value.strftime.return_value = '2025-05-14 14:54:24' - yield mock_dt - - -@pytest.fixture -def faker_instance(mock_kafka_producer): - logger = MagicMock() - notification_handler = MagicMock(spec=NotificationHandler) - return Faker( - bootstrap_servers='localhost:9092', - logger=logger, - notification_handler=notification_handler - ) - - -@pytest.mark.asyncio -async def test_faker_init(faker_instance, mock_kafka_producer): - """Test Faker initialization with correct parameters""" - assert faker_instance.producer is not None - assert len(faker_instance.tags) == 6 - -metadata = { - 'metadata': { - 'model_id': 'test_model_id', - 'model_name': 'test_model', - 'schedule_name': 'test_schedule', - 'workflow_name': 'scouter' - } -} - - -@pytest.mark.asyncio -async def test_generate_and_send_data_default_count(faker_instance, - mock_kafka_producer, mock_datetime): - """Test generating data with default message count""" - # Mock random.choice to control the output - with patch('random.choice') as mock_choice, \ - patch('random.uniform', return_value=42.5), \ - patch('random.randint', return_value=3), \ - patch('random.random', return_value=0.5): - - # Setup mock for tag and name selection - mock_choice.side_effect = [ - 'ns=1;i=1001', - 'ns=1;i=1002', - 'ns=1;i=1003' - ] - - # Call the method - await faker_instance.generate_and_send_data({'topic': 'test_topic', **metadata}) - - # Verify the producer was called 3 times (default count) - assert mock_kafka_producer.send.call_count == 3 - mock_kafka_producer.flush.assert_called_once() - - # Verify the message format - expected_data = [{ - 'tag': 'ns=1;i=1001', - 'name': 'Temperature Sensor', - 'timestamp': '2025-05-14 14:54:24', - 'value': 42.5 - }, { - 'tag': 'ns=1;i=1002', - 'name': 'Vibration Meter', - 'timestamp': '2025-05-14 14:54:24', - 'value': 42.5 - }, { - 'tag': 'ns=1;i=1003', - 'name': 'Pressure Gauge', - 'timestamp': '2025-05-14 14:54:24', - 'value': 42.5 - }] - mock_kafka_producer.send.assert_has_calls([ - call('test_topic', value=expected_data[0]), - call('test_topic', value=expected_data[1]), - call('test_topic', value=expected_data[2]) - ]) - - -@pytest.mark.asyncio -async def test_generate_and_send_data_custom_count(faker_instance, mock_kafka_producer): - """Test generating data with custom message count""" - # Call the method with custom count - await faker_instance.generate_and_send_data({ - 'topic': 'test_topic', - 'num_messages': 2, - **metadata - }) - - # Verify the producer was called 2 times - assert mock_kafka_producer.send.call_count == 2 - mock_kafka_producer.flush.assert_called_once() - - -@pytest.mark.asyncio -async def test_generate_and_send_data_no_topic(faker_instance): - """Test that ValueError is raised when no topic is provided""" - with pytest.raises(ValueError, match="Topic must be specified in input_data"): - await faker_instance.generate_and_send_data({**metadata}) - - -@pytest.mark.asyncio -@patch('scouter.activities.faker.random.random', return_value=0.5) -async def test_generate_and_send_data_random_values(_random, faker_instance, mock_kafka_producer): - """Test that random values are within expected ranges""" - # Call the method - await faker_instance.generate_and_send_data({'topic': 'test_topic', **metadata}) - - # Get the call arguments - call_args = mock_kafka_producer.send.call_args[1]['value'] - - # Verify the data structure - assert 'tag' in call_args - assert call_args['tag'] in faker_instance.tags - assert 'value' in call_args - assert 0 <= call_args['value'] <= 100 - - -@pytest.mark.asyncio -@patch('scouter.activities.faker.random.random', return_value=0.05) -async def test_generate_and_send_data_generate_null_values( - _random_mock, - faker_instance, - mock_kafka_producer): - # Call the method - await faker_instance.generate_and_send_data({'topic': 'test_topic', - 'num_messages': 1, - **metadata}) - - # Get the call arguments - call_args = mock_kafka_producer.send.call_args[1]['value'] - - assert call_args['value'] is None - assert call_args['tag'] in faker_instance.tags - assert 'name' in call_args - assert 'timestamp' in call_args diff --git a/tests/activities/test_gates.py b/tests/activities/test_gates.py index f0057fd..7f0f026 100644 --- a/tests/activities/test_gates.py +++ b/tests/activities/test_gates.py @@ -1,8 +1,11 @@ -from unittest.mock import Mock, patch, MagicMock, ANY +from typing import Any +from unittest.mock import ANY, MagicMock, Mock, patch + import numpy as np import pandas as pd import pytest from sientia_do.notifications.models import NotificationLevel + from scouter.activities.gates import Gates @@ -21,7 +24,7 @@ metadata = { 'model_id': 'test_model_id', 'model_name': 'test_model', 'schedule_name': 'test_schedule', - 'workflow_name': 'scouter' + 'workflow_name': 'scouter', } } @@ -31,11 +34,7 @@ async def test_data_quality_gate_with_null_values_filter_discard(gates_fixture): """Test data_quality_gate with NULL_VALUES_FILTER and DISCARD policy.""" # Setup test data input_data = { - 'filters': { - 'NULL_VALUES_FILTER': { - 'policy': 'DISCARD' - } - }, + 'filters': {'NULL_VALUES_FILTER': {'policy': 'DISCARD'}}, 'data': { 'name': ['tag1', 'tag2', 'tag3'], 'tag': ['tag1', 'tag2', 'tag3'], @@ -45,9 +44,9 @@ async def test_data_quality_gate_with_null_values_filter_discard(gates_fixture): 'model_tags': { 'tag1': {'data_range': [0, 100]}, 'tag2': {'data_range': [0, 100]}, - 'tag3': {'data_range': [0, 100]} + 'tag3': {'data_range': [0, 100]}, }, - **metadata + **metadata, } # Execute @@ -64,29 +63,26 @@ async def test_data_quality_gate_with_out_of_bounds_filter_keep(gates_fixture): """Test data_quality_gate with OUT_OF_BOUNDS_FILTER and KEEP policy.""" # Setup test data with out of bounds values input_data = { - 'filters': { - 'OUT_OF_BOUNDS_FILTER': { - 'policy': 'KEEP' - } - }, + 'filters': {'OUT_OF_BOUNDS_FILTER': {'policy': 'KEEP'}}, 'data': { 'name': ['tag1', 'tag2', 'tag3'], 'tag': ['tag1', 'tag2', 'tag3'], 'value': [1.0, 200.0, 3.0], - 'timestamp': ['2023-01-01', '2023-01-02', '2023-01-03'] + 'timestamp': ['2023-01-01', '2023-01-02', '2023-01-03'], }, 'model_tags': { 'tag1': {'data_range': [0, 100]}, 'tag2': {'data_range': [0, 100]}, - 'tag3': {'data_range': [0, 100]} + 'tag3': {'data_range': [0, 100]}, }, - **metadata + **metadata, } # Mock the out_of_bounds_filter to return rows with out of bounds values - with patch('scouter.activities.gates.quality_gate_filters', { - 'OUT_OF_BOUNDS_FILTER': lambda df: df[df['tag'] == 'tag2'] - }): + with patch( + 'scouter.activities.gates.quality_gate_filters', + {'OUT_OF_BOUNDS_FILTER': lambda df: df[df['tag'] == 'tag2']}, + ): # Execute result = await gates_fixture.data_quality_gate(input_data) @@ -101,33 +97,33 @@ async def test_data_quality_gate_with_multiple_filters(gates_fixture): # Setup test data input_data = { 'filters': { - 'NULL_VALUES_FILTER': { - 'policy': 'DISCARD' - }, - 'OUT_OF_BOUNDS_FILTER': { - 'policy': 'DISCARD' - } + 'NULL_VALUES_FILTER': {'policy': 'DISCARD'}, + 'OUT_OF_BOUNDS_FILTER': {'policy': 'DISCARD'}, }, 'data': { 'tag': ['tag1', 'tag2', 'tag3', 'tag4'], 'name': ['tag1', 'tag2', 'tag3', 'tag4'], 'value': [1.0, None, 300.0, 4.0], - 'timestamp': ['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-04'] + 'timestamp': ['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-04'], }, 'model_tags': { 'tag1': {'data_range': [0, 100]}, 'tag2': {'data_range': [0, 100]}, 'tag3': {'data_range': [0, 100]}, - 'tag4': {'data_range': [0, 100]} + 'tag4': {'data_range': [0, 100]}, }, - **metadata + **metadata, } result = await gates_fixture.data_quality_gate(input_data) # Verify only tag1 and tag4 remain (tag2 has null, tag3 is out of bounds) - assert result == {'tag': {0: 'tag1', 3: 'tag4'}, 'name': {0: 'tag1', 3: 'tag4'}, 'value': { - 0: 1.0, 3: 4.0}, 'timestamp': {0: '2023-01-01', 3: '2023-01-04'}} + assert result == { + 'tag': {0: 'tag1', 3: 'tag4'}, + 'name': {0: 'tag1', 3: 'tag4'}, + 'value': {0: 1.0, 3: 4.0}, + 'timestamp': {0: '2023-01-01', 3: '2023-01-04'}, + } # Should be called twice (once for each filter) assert gates_fixture.send_notification.call_count == 2 @@ -138,21 +134,10 @@ async def test_data_quality_gate_with_unknown_filter(gates_fixture): # Setup test data with unknown filter gates_fixture.warning = MagicMock() input_data = { - 'filters': { - 'UNKNOWN_FILTER': { - 'policy': 'DISCARD' - } - }, - 'data': { - 'tag': ['tag1'], - 'name': ['tag1'], - 'value': [1.0], - 'timestamp': ['2023-01-01'] - }, - 'model_tags': { - 'tag1': {'data_range': [0, 100]} - }, - **metadata + 'filters': {'UNKNOWN_FILTER': {'policy': 'DISCARD'}}, + 'data': {'tag': ['tag1'], 'name': ['tag1'], 'value': [1.0], 'timestamp': ['2023-01-01']}, + 'model_tags': {'tag1': {'data_range': [0, 100]}}, + **metadata, } # Execute @@ -161,8 +146,7 @@ async def test_data_quality_gate_with_unknown_filter(gates_fixture): # Verify data is unchanged and warning is logged assert len(result['tag']) == 1 gates_fixture.warning.assert_called_once_with( - "Filter UNKNOWN_FILTER not found", - metadata=metadata['metadata'] + 'Filter UNKNOWN_FILTER not found', metadata=metadata['metadata'] ) @@ -171,30 +155,19 @@ async def test_data_quality_gate_with_filter_error(gates_fixture): """Test data_quality_gate when a filter raises an exception.""" # Setup test data input_data = { - 'filters': { - 'NULL_VALUES_FILTER': { - 'policy': 'DISCARD' - } - }, - 'data': { - 'tag': ['tag1'], - 'name': ['tag1'], - 'value': [1.0], - 'timestamp': ['2023-01-01'] - }, - 'model_tags': { - 'tag1': {'data_range': [0, 100]} - }, - **metadata + 'filters': {'NULL_VALUES_FILTER': {'policy': 'DISCARD'}}, + 'data': {'tag': ['tag1'], 'name': ['tag1'], 'value': [1.0], 'timestamp': ['2023-01-01']}, + 'model_tags': {'tag1': {'data_range': [0, 100]}}, + **metadata, } # Mock the filter to raise an exception def failing_filter(_, _model_tags): - raise ValueError("Filter error") + raise ValueError('Filter error') - with patch('scouter.activities.gates.quality_gate_filters', { - 'NULL_VALUES_FILTER': failing_filter - }): + with patch( + 'scouter.activities.gates.quality_gate_filters', {'NULL_VALUES_FILTER': failing_filter} + ): # Execute result = await gates_fixture.data_quality_gate(input_data) @@ -202,9 +175,9 @@ async def test_data_quality_gate_with_filter_error(gates_fixture): assert len(result['tag']) == 1 gates_fixture.send_notification.assert_called_once() call_args = gates_fixture.send_notification.call_args[1] - assert call_args['notification_id'] == "DATA_QUALITY_GATE_ISSUES" + assert call_args['notification_id'] == 'DATA_QUALITY_GATE_ISSUES' assert call_args['level'] == NotificationLevel.ERROR - assert "Filter error" in call_args['message'] + assert 'Filter error' in call_args['message'] @pytest.mark.asyncio @@ -212,19 +185,10 @@ async def test_data_quality_gate_with_empty_data(gates_fixture): """Test data_quality_gate with empty input data.""" # Setup empty input data input_data = { - 'filters': { - 'NULL_VALUES_FILTER': { - 'policy': 'DISCARD' - } - }, - 'data': { - 'tag': [], - 'name': [], - 'value': [], - 'timestamp': [] - }, + 'filters': {'NULL_VALUES_FILTER': {'policy': 'DISCARD'}}, + 'data': {'tag': [], 'name': [], 'value': [], 'timestamp': []}, 'model_tags': {}, - **metadata + **metadata, } # Execute @@ -241,16 +205,9 @@ async def test_data_quality_gate_with_no_filters(gates_fixture): # Setup test data with no filters input_data = { 'filters': {}, - 'data': { - 'tag': ['tag1'], - 'name': ['tag1'], - 'value': [1.0], - 'timestamp': ['2023-01-01'] - }, - 'model_tags': { - 'tag1': {'data_range': [0, 100]} - }, - **metadata + 'data': {'tag': ['tag1'], 'name': ['tag1'], 'value': [1.0], 'timestamp': ['2023-01-01']}, + 'model_tags': {'tag1': {'data_range': [0, 100]}}, + **metadata, } # Execute @@ -262,7 +219,7 @@ async def test_data_quality_gate_with_no_filters(gates_fixture): @pytest.mark.parametrize( - "group_data, aggr_function, expected_result", + 'group_data, aggr_function, expected_result', [ # Single value case (pd.DataFrame({'value': [10.0]}), 'avg', 10.0), @@ -273,18 +230,16 @@ async def test_data_quality_gate_with_no_filters(gates_fixture): (pd.DataFrame({'value': [1.0, 2.0, 3.0, 4.0]}), 'min', 1.0), (pd.DataFrame({'value': [1.0, 2.0, 3.0, 4.0]}), 'lts', 4.0), # With NaN values - (pd.DataFrame({'value': [1.0, np.nan, 3.0, 4.0]}), - 'avg', 2.6666666666666665), + (pd.DataFrame({'value': [1.0, np.nan, 3.0, 4.0]}), 'avg', 2.6666666666666665), # Empty group after dropping NaN (pd.DataFrame({'value': [np.nan, np.nan]}), 'avg', None), # Invalid aggregation function (pd.DataFrame({'value': [1.0, 2.0]}), 'invalid', 'continue'), - ] + ], ) def test_apply_aggregation(gates_fixture, group_data, aggr_function, expected_result): """Test apply_aggregation method with various scenarios.""" - result = gates_fixture.apply_aggregation( - group_data, aggr_function, metadata) + result = gates_fixture.apply_aggregation(group_data, aggr_function, metadata) assert result == expected_result # Check notification was sent for invalid function @@ -305,22 +260,23 @@ async def test_aggregate_data(gates_fixture): {'tag': 'tag2', 'name': 'name2', 'value': 4.0, 'timestamp': '2023-01-01'}, {'tag': 'tag2', 'name': 'name2', 'value': 5.0, 'timestamp': '2023-01-02'}, {'tag': 'tag2', 'name': 'name2', 'value': 6.0, 'timestamp': '2023-01-03'}, - {'tag': 'tag1', 'name': 'name1', - 'value': None, 'timestamp': '2023-01-04'}, + {'tag': 'tag1', 'name': 'name1', 'value': None, 'timestamp': '2023-01-04'}, ], 'model_tags': { 'name1': {'aggr_func': 'avg'}, 'name2': {'aggr_func': 'max'}, }, - **metadata + **metadata, } # Expected result - expected_result = {'tag': {0: 'tag1', 1: 'tag2'}, - 'name': {0: 'name1', 1: 'name2'}, - 'value': {0: 2.0, 1: 6.0}, - 'timestamp': {0: '2023-01-04', 1: '2023-01-03'}, - 'aggregation_function': {0: 'avg', 1: 'max'}} + expected_result = { + 'tag': {0: 'tag1', 1: 'tag2'}, + 'name': {0: 'name1', 1: 'name2'}, + 'value': {0: 2.0, 1: 6.0}, + 'timestamp': {0: '2023-01-04', 1: '2023-01-03'}, + 'aggregation_function': {0: 'avg', 1: 'max'}, + } # Execute result = await gates_fixture.aggregate_data(input_data) @@ -332,7 +288,6 @@ async def test_aggregate_data(gates_fixture): @pytest.mark.asyncio async def test_aggregate_data_with_continue(gates_fixture): - gates_fixture.apply_aggregation = MagicMock(return_value='continue') input_data = { @@ -343,18 +298,17 @@ async def test_aggregate_data_with_continue(gates_fixture): {'tag': 'tag2', 'name': 'name2', 'value': 4.0, 'timestamp': '2023-01-01'}, {'tag': 'tag2', 'name': 'name2', 'value': 5.0, 'timestamp': '2023-01-02'}, {'tag': 'tag2', 'name': 'name2', 'value': 6.0, 'timestamp': '2023-01-03'}, - {'tag': 'tag1', 'name': 'name1', - 'value': None, 'timestamp': '2023-01-04'}, + {'tag': 'tag1', 'name': 'name1', 'value': None, 'timestamp': '2023-01-04'}, ], 'model_tags': { 'name1': {'aggr_function': 'avg'}, 'name2': {'aggr_function': 'max'}, }, - **metadata + **metadata, } # Expected result - expected_result = {} + expected_result: dict[str, Any] = {} # Execute result = await gates_fixture.aggregate_data(input_data) @@ -366,9 +320,7 @@ async def test_aggregate_data_with_continue(gates_fixture): @pytest.mark.asyncio async def test_aggregate_data_raise_exception(gates_fixture): - - gates_fixture.apply_aggregation = MagicMock( - side_effect=Exception("Test exception")) + gates_fixture.apply_aggregation = MagicMock(side_effect=Exception('Test exception')) input_data = { 'data': [ @@ -378,42 +330,39 @@ async def test_aggregate_data_raise_exception(gates_fixture): {'tag': 'tag2', 'name': 'name2', 'value': 4.0, 'timestamp': '2023-01-01'}, {'tag': 'tag2', 'name': 'name2', 'value': 5.0, 'timestamp': '2023-01-02'}, {'tag': 'tag2', 'name': 'name2', 'value': 6.0, 'timestamp': '2023-01-03'}, - {'tag': 'tag1', 'name': 'name1', - 'value': None, 'timestamp': '2023-01-04'}, + {'tag': 'tag1', 'name': 'name1', 'value': None, 'timestamp': '2023-01-04'}, ], 'model_tags': { 'name1': {'aggr_function': 'avg'}, 'name2': {'aggr_function': 'max'}, }, - **metadata + **metadata, } try: await gates_fixture.aggregate_data(input_data) except Exception as e: - assert str(e) == "Test exception" + assert str(e) == 'Test exception' gates_fixture.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="AGGREGATION_ISSUES", - message="Error aggregating data: Test exception", - block="aggregate_data", + notification_id='AGGREGATION_ISSUES', + message='Error aggregating data: Test exception', + block='aggregate_data', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False + raise AssertionError('Exception not raised') @pytest.mark.asyncio @patch('scouter.activities.gates.metrics') async def test_write_metrics(mock_metrics, gates_fixture): """Test write_metrics method.""" - input_data = { - 'metadata': metadata['metadata'] - } + input_data = {'metadata': metadata['metadata']} await gates_fixture.write_metrics(input_data) mock_metrics.LABORIOUS_DATA_WRITTEN_COUNT.labels.assert_called_once_with( pod_id=gates_fixture.pod_id, model_name=metadata['metadata']['model_name'], - pipeline_name=metadata['metadata']['workflow_name'] + pipeline_name=metadata['metadata']['workflow_name'], ) diff --git a/tests/activities/test_mongo.py b/tests/activities/test_mongo.py index 5f91e1c..0d62e0a 100644 --- a/tests/activities/test_mongo.py +++ b/tests/activities/test_mongo.py @@ -1,8 +1,10 @@ from datetime import datetime -from unittest.mock import ANY, MagicMock, call, patch +from unittest.mock import ANY, MagicMock, patch + from pytest import fixture, mark from sientia_do.notifications.models import NotificationLevel from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ + from scouter.activities.mongodb import MongoDB, clear_mongo_id @@ -10,15 +12,9 @@ def test_clear_mongo_id(): """Test clear_mongo_id""" data = [ {'_id': '1', 'name': 'test1'}, - {'_id': '2', 'name': [{ - '_id': '3', - 'name': 'test3' - }]}, - {'_id': '4', 'name': { - '_id': '5', - 'name': 'test2' - }}, - [{'_id': '6', 'name': 'test2'}] + {'_id': '2', 'name': [{'_id': '3', 'name': 'test3'}]}, + {'_id': '4', 'name': {'_id': '5', 'name': 'test2'}}, + [{'_id': '6', 'name': 'test2'}], ] result = clear_mongo_id(data) @@ -27,7 +23,7 @@ def test_clear_mongo_id(): {'name': 'test1'}, {'name': [{'name': 'test3'}]}, {'name': {'name': 'test2'}}, - [{'name': 'test2'}] + [{'name': 'test2'}], ] @@ -38,18 +34,16 @@ def test_mongodb___init__(mock_mongo_client): connection_string='mongodb://localhost:27017', database_name='test_db', logger=MagicMock(), - notification_handler=MagicMock() + notification_handler=MagicMock(), ) mock_mongo_client.assert_called_once_with( - 'mongodb://localhost:27017', - serverSelectionTimeoutMS=5000 + 'mongodb://localhost:27017', serverSelectionTimeoutMS=5000 ) mock_mongo_client.return_value.server_info.assert_called_once() - mock_mongo_client.return_value.__getitem__.assert_called_once_with( - 'test_db') + mock_mongo_client.return_value.__getitem__.assert_called_once_with('test_db') assert mongo.client is not None assert mongo.database is not None @@ -63,7 +57,7 @@ def mongodb_activity(mock_mongo_client): connection_string='mongodb://localhost:27017', database_name='test_db', logger=MagicMock(), - notification_handler=MagicMock() + notification_handler=MagicMock(), ) return mongo @@ -96,34 +90,27 @@ async def test_load_latest_data_none_last_data_timestamp(mongodb_activity): 'name': 'test1', 'value': 1, 'inserted_at': datetime.strptime( - '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ) + '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ + ), } ] - result = await mongodb_activity.load_latest_data({ - 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, - 'collection_name': 'test_collection', - 'last_data_timestamp': None - }) - - mongodb_activity.database.__getitem__.assert_called_once_with( - 'test_collection') - - collection.find.assert_called_once_with( - {}, - {"_id": 0} + result = await mongodb_activity.load_latest_data( + { + 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, + 'collection_name': 'test_collection', + 'last_data_timestamp': None, + } ) + mongodb_activity.database.__getitem__.assert_called_once_with('test_collection') + + collection.find.assert_called_once_with({}, {'_id': 0}) + assert result == { - 'name': { - 0: 'test1' - }, - 'value': { - 0: 1 - }, - 'inserted_at': { - 0: '2023-01-01 12:00:00.000000+0000' - } + 'name': {0: 'test1'}, + 'value': {0: 1}, + 'inserted_at': {0: '2023-01-01 12:00:00.000000+0000'}, } @@ -138,39 +125,36 @@ async def test_load_latest_data_not_none_last_data_timestamp(mongodb_activity): 'name': 'test1', 'value': 1, 'inserted_at': datetime.strptime( - '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ) + '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ + ), } ] - result = await mongodb_activity.load_latest_data({ - 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, - 'collection_name': 'test_collection', - 'last_data_timestamp': '2023-01-01 12:00:00.000000+0000' - }) + result = await mongodb_activity.load_latest_data( + { + 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, + 'collection_name': 'test_collection', + 'last_data_timestamp': '2023-01-01 12:00:00.000000+0000', + } + ) - mongodb_activity.database.__getitem__.assert_called_once_with( - 'test_collection') + mongodb_activity.database.__getitem__.assert_called_once_with('test_collection') collection.find.assert_called_once_with( { 'inserted_at': { '$gt': datetime.strptime( - '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ) + '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ + ) } }, - {"_id": 0} + {'_id': 0}, ) assert result == { - 'name': { - 0: 'test1' - }, - 'value': { - 0: 1 - }, - 'inserted_at': { - 0: '2023-01-01 12:00:00.000000+0000' - } + 'name': {0: 'test1'}, + 'value': {0: 1}, + 'inserted_at': {0: '2023-01-01 12:00:00.000000+0000'}, } @@ -184,20 +168,21 @@ async def test_load_latest_data_error(mongodb_activity): collection.find.side_effect = Exception('test') try: - await mongodb_activity.load_latest_data({ - 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, - 'collection_name': 'test_collection', - 'last_data_timestamp': '2023-01-01 12:00:00.000000+0000' - }) + await mongodb_activity.load_latest_data( + { + 'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, + 'collection_name': 'test_collection', + 'last_data_timestamp': '2023-01-01 12:00:00.000000+0000', + } + ) except Exception as e: assert str(e) == 'test' mongodb_activity.send_notification.assert_called_once_with( - metadata={'workflow_name': 'test_pipeline', - 'schedule_name': 'test_schedule'}, + metadata={'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, notification_id='MONGO_LOAD_ERROR', message='Error loading data from MongoDB: test', block='load_latest_data', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) diff --git a/tests/activities/test_redis.py b/tests/activities/test_redis.py index f87d46f..2340ed6 100644 --- a/tests/activities/test_redis.py +++ b/tests/activities/test_redis.py @@ -1,10 +1,12 @@ -from unittest.mock import MagicMock, patch, ANY from datetime import datetime -import pytest +from unittest.mock import ANY, MagicMock, patch + import numpy as np +import pytest from pandas import DataFrame from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.models import NotificationLevel + from scouter.activities.redis import Redis @@ -13,9 +15,14 @@ from scouter.activities.redis import Redis def redis_activity(_mock_redis_init): logger = MagicMock() notification_handler = MagicMock(spec=NotificationHandler) - activity = Redis(host='localhost', port=6379, - logger=logger, notification_handler=notification_handler, - username='test', password='test') + activity = Redis( + host='localhost', + port=6379, + logger=logger, + notification_handler=notification_handler, + username='test', + password='test', + ) activity.redis_client = MagicMock() activity.logger = logger @@ -29,17 +36,16 @@ def test_redis_initialization(mock_redis_init): """Test Redis activity initialization""" logger = MagicMock() notification_handler = MagicMock(spec=NotificationHandler) - Redis(host='localhost', port=6379, - logger=logger, notification_handler=notification_handler, - username='test', password='test') + Redis( + host='localhost', + port=6379, + logger=logger, + notification_handler=notification_handler, + username='test', + password='test', + ) mock_redis_init.assert_called_once_with( - ANY, - 'localhost', - 6379, - 'test', - 'test', - logger, - notification_handler + ANY, 'localhost', 6379, 'test', 'test', logger, notification_handler ) @@ -48,7 +54,7 @@ metadata = { 'model_id': 'test_model_id', 'model_name': 'test_model', 'schedule_name': 'test_schedule', - 'workflow_name': 'scouter' + 'workflow_name': 'scouter', } } @@ -56,11 +62,7 @@ metadata = { @pytest.mark.asyncio async def test_get_last_data_timestamp_none(redis_activity): """Test get_last_data_timestamp""" - test_data = { - **metadata, - 'workflow_name': 'test_pipeline', - 'schedule_name': 'test_schedule' - } + test_data = {**metadata, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'} redis_activity.get = MagicMock(return_value=None) @@ -72,19 +74,13 @@ async def test_get_last_data_timestamp_none(redis_activity): @pytest.mark.asyncio async def test_get_last_data_timestamp_not_none(redis_activity): """Test get_last_data_timestamp""" - test_data = { - **metadata, - 'workflow_name': 'test_pipeline', - 'schedule_name': 'test_schedule' - } + test_data = {**metadata, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'} redis_activity.get = MagicMock(return_value='2023-01-01 12:00:00') result = await redis_activity.get_last_data_timestamp(test_data) - redis_activity.get.assert_called_once_with( - 'last_data_timestamp:test_pipeline:test_schedule' - ) + redis_activity.get.assert_called_once_with('last_data_timestamp:test_pipeline:test_schedule') assert result == '2023-01-01 12:00:00' @@ -92,17 +88,12 @@ async def test_get_last_data_timestamp_not_none(redis_activity): @pytest.mark.asyncio async def test_get_last_data_timestamp_error(redis_activity): """Test get_last_data_timestamp error""" - test_data = { - **metadata, - 'workflow_name': 'test_pipeline', - 'schedule_name': 'test_schedule' - } + test_data = {**metadata, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'} redis_activity.send_notification = MagicMock() redis_activity.get = MagicMock(side_effect=Exception('test')) try: - await redis_activity.get_last_data_timestamp(test_data) except Exception as e: @@ -110,15 +101,15 @@ async def test_get_last_data_timestamp_error(redis_activity): redis_activity.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="REDIS_GET_ERROR", - message="Error getting last data timestamp: test", - block="get_last_data_timestamp", + notification_id='REDIS_GET_ERROR', + message='Error getting last data timestamp: test', + block='get_last_data_timestamp', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "Expected exception" + raise AssertionError('Expected exception') @pytest.mark.asyncio @@ -128,7 +119,7 @@ async def test_put_last_data_timestamp_empty_dataframe(redis_activity): **metadata, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule', - 'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records') + 'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records'), } redis_activity.set = MagicMock() @@ -144,16 +135,18 @@ async def test_put_last_data_timestamp_empty_dataframe(redis_activity): async def test_put_last_data_timestamp_not_empty_dataframe(redis_activity): """Test put_last_data_timestamp with not empty dataframe""" - data = DataFrame({ - 'name': ['sensor1', 'sensor2'], - 'value': [25.5, 30.0], - 'inserted_at': ['2023-01-01 12:00:00', '2023-01-01 12:00:01'] - }) + data = DataFrame( + { + 'name': ['sensor1', 'sensor2'], + 'value': [25.5, 30.0], + 'inserted_at': ['2023-01-01 12:00:00', '2023-01-01 12:00:01'], + } + ) test_data = { **metadata, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule', - 'data': data.to_dict('records') + 'data': data.to_dict('records'), } redis_activity.set = MagicMock() @@ -163,9 +156,7 @@ async def test_put_last_data_timestamp_not_empty_dataframe(redis_activity): assert result == '2023-01-01 12:00:01' redis_activity.set.assert_called_once_with( - 'last_data_timestamp:test_pipeline:test_schedule', - '2023-01-01 12:00:01', - ttl=18000 + 'last_data_timestamp:test_pipeline:test_schedule', '2023-01-01 12:00:01', ttl=18000 ) @@ -176,11 +167,13 @@ async def test_put_last_data_timestamp_error(redis_activity): **metadata, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule', - 'data': DataFrame({ - 'name': ['sensor1', 'sensor2'], - 'value': [25.5, 30.0], - 'inserted_at': ['2023-01-01 12:00:00'] * 2 - }).to_dict('records') + 'data': DataFrame( + { + 'name': ['sensor1', 'sensor2'], + 'value': [25.5, 30.0], + 'inserted_at': ['2023-01-01 12:00:00'] * 2, + } + ).to_dict('records'), } redis_activity.send_notification = MagicMock() @@ -194,15 +187,15 @@ async def test_put_last_data_timestamp_error(redis_activity): redis_activity.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="REDIS_SET_ERROR", - message="Error setting last data timestamp: test", - block="put_last_data_timestamp", + notification_id='REDIS_SET_ERROR', + message='Error setting last data timestamp: test', + block='put_last_data_timestamp', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "Expected exception" + raise AssertionError('Expected exception') @pytest.mark.asyncio @@ -215,15 +208,14 @@ async def test_group_and_hold_data_new_key(redis_activity): 'schedule_name': 'test_schedule', 'retention_time': 3600, 'model_id': 1, - 'data': DataFrame({ - 'name': ['sensor1', 'sensor2'], - 'value': [25.5, 30.0], - 'timestamp': ['2023-01-01 12:00:00'] * 2 - }).to_dict('records'), - 'model_tags': { - 'sensor1': 'sensor1', - 'sensor2': 'sensor2' - } + 'data': DataFrame( + { + 'name': ['sensor1', 'sensor2'], + 'value': [25.5, 30.0], + 'timestamp': ['2023-01-01 12:00:00'] * 2, + } + ).to_dict('records'), + 'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'}, } # Mock get to return None for new key @@ -238,7 +230,7 @@ async def test_group_and_hold_data_new_key(redis_activity): 'timestamp': {0: '2023-01-01 12:00:00', 1: '2023-01-01 12:00:00'}, 'variable': {0: 'sensor1', 1: 'sensor2'}, 'value': {0: 25.5, 1: 30.0}, - 'model_id': {0: 1, 1: 1} + 'model_id': {0: 1, 1: 1}, } assert result == expected_result @@ -246,11 +238,7 @@ async def test_group_and_hold_data_new_key(redis_activity): redis_activity.set.assert_called_once() args, kwargs = redis_activity.set.call_args assert args[0] == 'held_data_test_pipeline_test_schedule' - assert args[1] == { - 'sensor1': 25.5, - 'sensor2': 30.0, - 'timestamp': '2023-01-01 12:00:00' - } + assert args[1] == {'sensor1': 25.5, 'sensor2': 30.0, 'timestamp': '2023-01-01 12:00:00'} assert kwargs['ttl'] == 3600 @@ -258,11 +246,7 @@ async def test_group_and_hold_data_new_key(redis_activity): async def test_group_and_hold_data_update_existing(redis_activity): """Test updating existing data with group_and_hold_data""" # Setup initial data in Redis - existing_data = { - 'sensor1': 20.0, - 'sensor2': 28.0, - 'timestamp': '2023-01-01 11:00:00' - } + existing_data = {'sensor1': 20.0, 'sensor2': 28.0, 'timestamp': '2023-01-01 11:00:00'} # New data to update with test_data = { @@ -271,16 +255,14 @@ async def test_group_and_hold_data_update_existing(redis_activity): 'schedule_name': 'test_schedule', 'retention_time': 3600, 'model_id': 1, - 'data': DataFrame({ - 'name': ['sensor1', 'sensor3'], - 'value': [25.5, 42.0], - 'timestamp': ['2023-01-01 12:00:00'] * 2 - }).to_dict('records'), - 'model_tags': { - 'sensor1': 'sensor1', - 'sensor2': 'sensor2', - 'sensor3': 'sensor3' - } + 'data': DataFrame( + { + 'name': ['sensor1', 'sensor3'], + 'value': [25.5, 42.0], + 'timestamp': ['2023-01-01 12:00:00'] * 2, + } + ).to_dict('records'), + 'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2', 'sensor3': 'sensor3'}, } # Mock get to return existing data @@ -295,7 +277,7 @@ async def test_group_and_hold_data_update_existing(redis_activity): 'timestamp': {0: '2023-01-01 12:00:00', 1: '2023-01-01 12:00:00', 2: '2023-01-01 12:00:00'}, 'variable': {0: 'sensor1', 1: 'sensor2', 2: 'sensor3'}, 'value': {0: 25.5, 1: 28.0, 2: 42.0}, - 'model_id': {0: 1, 1: 1, 2: 1} + 'model_id': {0: 1, 1: 1, 2: 1}, } assert result == expected_result @@ -307,7 +289,7 @@ async def test_group_and_hold_data_update_existing(redis_activity): 'sensor1': 25.5, 'sensor2': 28.0, 'sensor3': 42.0, - 'timestamp': '2023-01-01 12:00:00' + 'timestamp': '2023-01-01 12:00:00', } assert kwargs['ttl'] == 3600 @@ -322,15 +304,14 @@ async def test_group_and_hold_data_with_none_values(redis_activity): 'schedule_name': 'test_schedule', 'retention_time': 3600, 'model_id': 1, - 'data': DataFrame({ - 'name': ['sensor1', 'sensor2'], - 'value': [None, 30.0], - 'timestamp': [datetime(2023, 1, 1, 12, 0, 0)] * 2 - }).to_dict('records'), - 'model_tags': { - 'sensor1': 'sensor1', - 'sensor2': 'sensor2' - } + 'data': DataFrame( + { + 'name': ['sensor1', 'sensor2'], + 'value': [None, 30.0], + 'timestamp': [datetime(2023, 1, 1, 12, 0, 0)] * 2, + } + ).to_dict('records'), + 'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'}, } # Mock get to return None for new key @@ -355,10 +336,7 @@ async def test_group_and_hold_data_empty_dataframe(redis_activity): 'schedule_name': 'test_schedule', 'retention_time': 3600, 'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records'), - 'model_tags': { - 'sensor1': 'sensor1', - 'sensor2': 'sensor2' - } + 'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'}, } redis_activity.get = MagicMock(return_value=None) @@ -379,10 +357,7 @@ async def test_group_and_hold_data_error_get(redis_activity): 'retention_time': 3600, 'model_id': 1, 'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records'), - 'model_tags': { - 'sensor1': 'sensor1', - 'sensor2': 'sensor2' - } + 'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'}, } redis_activity.get = MagicMock(side_effect=Exception('test')) @@ -396,15 +371,15 @@ async def test_group_and_hold_data_error_get(redis_activity): redis_activity.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="REDIS_GET_ERROR", - message="Error getting held data: test", - block="group_and_hold_data", + notification_id='REDIS_GET_ERROR', + message='Error getting held data: test', + block='group_and_hold_data', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) else: - assert False, "Expected exception" + raise AssertionError('Expected exception') @pytest.mark.asyncio @@ -417,17 +392,10 @@ async def test_group_and_hold_data_error_set(redis_activity): 'retention_time': 3600, 'model_id': 1, 'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records'), - 'model_tags': { - 'sensor1': 'sensor1', - 'sensor2': 'sensor2' - } + 'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'}, } - existing_data = { - 'sensor1': 20.0, - 'sensor2': 28.0, - 'timestamp': '2023-01-01 11:00:00' - } + existing_data = {'sensor1': 20.0, 'sensor2': 28.0, 'timestamp': '2023-01-01 11:00:00'} # Mock get to return existing data redis_activity.get = MagicMock(return_value=existing_data) @@ -450,67 +418,65 @@ async def test_store_data_package(redis_activity): **metadata, 'workflow_name': 'test_workflow', 'schedule_name': 'test_schedule', - 'held_data': DataFrame({ - 'name': ['sensor1', 'sensor2'], - 'value': [25.5, 30.0], - 'timestamp': ['2023-01-01 12:00:00'] * 2 - }).to_dict(), - 'data': DataFrame({ - 'name': ['sensor1', 'sensor2'], - 'value': [25.5, 30.0], - 'timestamp': ['2023-01-01 12:00:00'] * 2 - }).to_dict(), - 'model_tags': { - 'sensor1': 'sensor1', - 'sensor2': 'sensor2' - } + 'held_data': DataFrame( + { + 'name': ['sensor1', 'sensor2'], + 'value': [25.5, 30.0], + 'timestamp': ['2023-01-01 12:00:00'] * 2, + } + ).to_dict(), + 'data': DataFrame( + { + 'name': ['sensor1', 'sensor2'], + 'value': [25.5, 30.0], + 'timestamp': ['2023-01-01 12:00:00'] * 2, + } + ).to_dict(), + 'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'}, } await redis_activity.store_data_package(test_data) redis_activity.set.assert_called_once_with( - ANY, - { - 'data': test_data['data'], - 'held_data': test_data['held_data'] - }, - ttl=120) + ANY, {'data': test_data['data'], 'held_data': test_data['held_data']}, ttl=120 + ) @pytest.mark.asyncio async def test_store_data_package_error(redis_activity): """Test store_data_package error""" - redis_activity.set = MagicMock(side_effect=Exception('test')) + redis_activity.set = MagicMock(side_effect=ValueError('test')) redis_activity.send_notification = MagicMock() test_data = { **metadata, 'workflow_name': 'test_workflow', 'schedule_name': 'test_schedule', - 'held_data': DataFrame({ - 'name': ['sensor1', 'sensor2'], - 'value': [25.5, 30.0], - 'timestamp': ['2023-01-01 12:00:00'] * 2 - }).to_dict(), - 'data': DataFrame({ - 'name': ['sensor1', 'sensor2'], - 'value': [25.5, 30.0], - 'timestamp': ['2023-01-01 12:00:00'] * 2 - }).to_dict(), - 'model_tags': { - 'sensor1': 'sensor1', - 'sensor2': 'sensor2' - } + 'held_data': DataFrame( + { + 'name': ['sensor1', 'sensor2'], + 'value': [25.5, 30.0], + 'timestamp': ['2023-01-01 12:00:00'] * 2, + } + ).to_dict(), + 'data': DataFrame( + { + 'name': ['sensor1', 'sensor2'], + 'value': [25.5, 30.0], + 'timestamp': ['2023-01-01 12:00:00'] * 2, + } + ).to_dict(), + 'model_tags': {'sensor1': 'sensor1', 'sensor2': 'sensor2'}, } - with pytest.raises(Exception): + with pytest.raises(ValueError): await redis_activity.store_data_package(test_data) redis_activity.send_notification.assert_called_once_with( metadata=metadata['metadata'], - notification_id="REDIS_SET_ERROR", - message="Error setting data package: test", - block="store_data_package", + notification_id='REDIS_SET_ERROR', + message='Error setting data package: test', + block='store_data_package', level=NotificationLevel.ERROR, - attachment_content=ANY + attachment_content=ANY, ) diff --git a/tests/test_metrics.py b/tests/test_metrics.py index 875113c..79a96d9 100644 --- a/tests/test_metrics.py +++ b/tests/test_metrics.py @@ -1,7 +1,7 @@ # tests/unit/test_metrics.py -import pytest -from prometheus_client import Counter, Gauge, Histogram +from prometheus_client import Counter, Gauge + import scouter.metrics as metrics # --- Test Functions for Each Metric (Corrected for v0.22.0 _name behavior) --- @@ -11,15 +11,22 @@ def test_scouter_laborious_data_written_count(): """Verify the definition of LABORIOUS_DATA_WRITTEN_COUNT.""" assert metrics.LABORIOUS_DATA_WRITTEN_COUNT is not None assert isinstance(metrics.LABORIOUS_DATA_WRITTEN_COUNT, Counter) - assert metrics.LABORIOUS_DATA_WRITTEN_COUNT._name == "scouter_laborious_data_written_count" + assert metrics.LABORIOUS_DATA_WRITTEN_COUNT._name == 'scouter_laborious_data_written_count' assert set(metrics.LABORIOUS_DATA_WRITTEN_COUNT._labelnames) == { - "pod_id", "model_name", "pipeline_name"} + 'pod_id', + 'model_name', + 'pipeline_name', + } def test_scouter_tag_changes_monitor(): """Verify the definition of TAG_CHANGES_MONITOR.""" assert metrics.TAG_CHANGES_MONITOR is not None assert isinstance(metrics.TAG_CHANGES_MONITOR, Gauge) - assert metrics.TAG_CHANGES_MONITOR._name == "scouter_tag_changes_monitor" + assert metrics.TAG_CHANGES_MONITOR._name == 'scouter_tag_changes_monitor' assert set(metrics.TAG_CHANGES_MONITOR._labelnames) == { - "pod_id", "model_name", "pipeline_name", "tag_name"} + 'pod_id', + 'model_name', + 'pipeline_name', + 'tag_name', + } diff --git a/tests/utils/quality/test_filters.py b/tests/utils/quality/test_filters.py index c08ab2c..7cc9cc5 100644 --- a/tests/utils/quality/test_filters.py +++ b/tests/utils/quality/test_filters.py @@ -1,8 +1,9 @@ -import pytest -import pandas as pd import numpy as np +import pandas as pd +import pytest from pandas.testing import assert_frame_equal -from scouter.utils.quality.filters import check_data_range, out_of_bounds_filter, null_values_filter + +from scouter.utils.quality.filters import check_data_range, null_values_filter, out_of_bounds_filter # Fixtures @@ -10,12 +11,14 @@ from scouter.utils.quality.filters import check_data_range, out_of_bounds_filter @pytest.fixture def sample_dataframe(): """Fixture providing a sample DataFrame for testing.""" - return pd.DataFrame({ - 'tag': ['temp', 'temp', 'pressure', 'pressure', 'humidity', 'wind_speed'], - 'name': ['temp', 'temp', 'pressure', 'pressure', 'humidity', 'wind_speed'], - 'value': [25, 35, 95, 105, 60, None], - 'timestamp': pd.date_range(start='2023-01-01', periods=6) - }) + return pd.DataFrame( + { + 'tag': ['temp', 'temp', 'pressure', 'pressure', 'humidity', 'wind_speed'], + 'name': ['temp', 'temp', 'pressure', 'pressure', 'humidity', 'wind_speed'], + 'value': [25, 35, 95, 105, 60, None], + 'timestamp': pd.date_range(start='2023-01-01', periods=6), + } + ) @pytest.fixture @@ -25,7 +28,7 @@ def nodes_data_range(): 'temp': {'data_range': [10, 30]}, 'pressure': {'data_range': [90, 100]}, 'humidity': {'data_range': [40, 80]}, - 'wind_speed': {'data_range': [0, 50]} + 'wind_speed': {'data_range': [0, 50]}, } @@ -59,6 +62,7 @@ def test_check_data_range(value, val_range, expected): else: assert result == expected + # Tests for out_of_bounds_filter @@ -72,8 +76,8 @@ def test_out_of_bounds_filter(sample_dataframe, nodes_data_range): 'timestamp': [ pd.Timestamp('2023-01-02'), pd.Timestamp('2023-01-04'), - pd.Timestamp('2023-01-06') - ] + pd.Timestamp('2023-01-06'), + ], } expected_df = pd.DataFrame(expected_data) @@ -91,6 +95,7 @@ def test_out_of_bounds_filter_empty_df(nodes_data_range): assert result.empty assert list(result.columns) == ['tag', 'name', 'value', 'timestamp'] + # Tests for null_values_filter @@ -100,7 +105,7 @@ def test_null_values_filter(sample_dataframe, nodes_data_range): 'tag': ['wind_speed'], 'name': ['wind_speed'], 'value': [None], - 'timestamp': [pd.Timestamp('2023-01-06')] + 'timestamp': [pd.Timestamp('2023-01-06')], } expected_df = pd.DataFrame(expected_data) @@ -113,12 +118,14 @@ def test_null_values_filter(sample_dataframe, nodes_data_range): def test_null_values_filter_no_nulls(nodes_data_range): """Test with a DataFrame containing no null values.""" - df = pd.DataFrame({ - 'tag': ['temp', 'pressure'], - 'name': ['temp', 'pressure'], - 'value': [25, 100], - 'timestamp': pd.date_range(start='2023-01-01', periods=2) - }) + df = pd.DataFrame( + { + 'tag': ['temp', 'pressure'], + 'name': ['temp', 'pressure'], + 'value': [25, 100], + 'timestamp': pd.date_range(start='2023-01-01', periods=2), + } + ) result = null_values_filter(df, nodes_data_range) assert result.empty assert list(result.columns) == ['tag', 'name', 'value', 'timestamp'] diff --git a/tests/utils/test_connectors_config.py b/tests/utils/test_connectors_config.py index e3c4bbd..61d2b0a 100644 --- a/tests/utils/test_connectors_config.py +++ b/tests/utils/test_connectors_config.py @@ -1,12 +1,14 @@ import os from unittest.mock import patch + import pytest + from scouter.utils.connectors_config import ( build_druid_config, + build_kafka_config, build_mongodb_config, build_postgres_config, - build_kafka_config, - build_redis_config + build_redis_config, ) @@ -16,7 +18,7 @@ def mock_env_vars(): yield -@pytest.mark.usefixtures("mock_env_vars") +@pytest.mark.usefixtures('mock_env_vars') def test_build_postgres_config_defaults(): """Test that build_postgres_config returns default values when no env vars are set""" config = build_postgres_config() @@ -28,22 +30,25 @@ def test_build_postgres_config_defaults(): 'password': 'sientia', 'dbname': 'sientia', 'min_connections': 5, - 'max_connections': 20 + 'max_connections': 20, } -@pytest.mark.usefixtures("mock_env_vars") +@pytest.mark.usefixtures('mock_env_vars') def test_build_postgres_config_with_env_vars(): """Test that build_postgres_config uses env vars when set""" - with patch.dict(os.environ, { - 'POSTGRES_HOST': 'db.example.com', - 'POSTGRES_PORT': '5433', - 'POSTGRES_USER': 'admin', - 'POSTGRES_PASSWORD': 'secret', - 'POSTGRES_DBNAME': 'test_db', - 'POSTGRES_MIN_CONNECTIONS': '3', - 'POSTGRES_MAX_CONNECTIONS': '15' - }): + with patch.dict( + os.environ, + { + 'POSTGRES_HOST': 'db.example.com', + 'POSTGRES_PORT': '5433', + 'POSTGRES_USER': 'admin', + 'POSTGRES_PASSWORD': 'secret', + 'POSTGRES_DBNAME': 'test_db', + 'POSTGRES_MIN_CONNECTIONS': '3', + 'POSTGRES_MAX_CONNECTIONS': '15', + }, + ): config = build_postgres_config() assert config == { @@ -53,11 +58,11 @@ def test_build_postgres_config_with_env_vars(): 'password': 'secret', 'dbname': 'test_db', 'min_connections': 3, - 'max_connections': 15 + 'max_connections': 15, } -@pytest.mark.usefixtures("mock_env_vars") +@pytest.mark.usefixtures('mock_env_vars') def test_build_kafka_config_defaults(): """Test that build_kafka_config returns default values when no env vars are set""" config = build_kafka_config() @@ -65,55 +70,53 @@ def test_build_kafka_config_defaults(): assert config == { 'bootstrap_servers': 'localhost:9092', 'polling_time': 1000, - 'group_id': 'scouter-group' + 'group_id': 'scouter-group', } -@pytest.mark.usefixtures("mock_env_vars") +@pytest.mark.usefixtures('mock_env_vars') def test_build_kafka_config_with_env_vars(): """Test that build_kafka_config uses env vars when set""" - with patch.dict(os.environ, { - 'KAFKA_BOOTSTRAP_SERVERS': 'kafka.example.com:9092', - 'KAFKA_POLLING_TIME': '5000' - }): + with patch.dict( + os.environ, + {'KAFKA_BOOTSTRAP_SERVERS': 'kafka.example.com:9092', 'KAFKA_POLLING_TIME': '5000'}, + ): config = build_kafka_config() assert config == { 'bootstrap_servers': 'kafka.example.com:9092', 'polling_time': 5000, - 'group_id': 'scouter-group' + 'group_id': 'scouter-group', } -@pytest.mark.usefixtures("mock_env_vars") +@pytest.mark.usefixtures('mock_env_vars') def test_build_redis_config_defaults(): """Test that build_redis_config returns default values when no env vars are set""" config = build_redis_config() - assert config == { - 'host': 'localhost', - 'port': 6379, - 'username': None, - 'password': None - } + assert config == {'host': 'localhost', 'port': 6379, 'username': None, 'password': None} -@pytest.mark.usefixtures("mock_env_vars") +@pytest.mark.usefixtures('mock_env_vars') def test_build_redis_config_with_env_vars(): """Test that build_redis_config uses env vars when set""" - with patch.dict(os.environ, { - 'REDIS_HOST': 'redis.example.com', - 'REDIS_PORT': '6380', - 'REDIS_USERNAME': 'test', - 'REDIS_PASSWORD': 'test' - }): + with patch.dict( + os.environ, + { + 'REDIS_HOST': 'redis.example.com', + 'REDIS_PORT': '6380', + 'REDIS_USERNAME': 'test', + 'REDIS_PASSWORD': 'test', + }, + ): config = build_redis_config() assert config == { 'host': 'redis.example.com', 'port': 6380, 'username': 'test', - 'password': 'test' + 'password': 'test', } @@ -123,23 +126,26 @@ def test_build_mongodb_config_defaults(): assert config == { 'connection_string': 'mongodb://sientia:sientia@localhost:27017', # NOSONAR - 'database_name': 'sientia' + 'database_name': 'sientia', } def test_build_mongodb_config_with_env_vars(): """Test that build_mongodb_config uses env vars when set""" - with patch.dict(os.environ, { - 'MONGODB_URL': 'mongodb.example.com:27017', - 'MONGODB_DATABASE_NAME': 'test_db', - 'MONGODB_USERNAME': 'test', - 'MONGODB_PASSWORD': 'test' - }): + with patch.dict( + os.environ, + { + 'MONGODB_URL': 'mongodb.example.com:27017', + 'MONGODB_DATABASE_NAME': 'test_db', + 'MONGODB_USERNAME': 'test', + 'MONGODB_PASSWORD': 'test', + }, + ): config = build_mongodb_config() assert config == { 'connection_string': 'mongodb://test:test@mongodb.example.com:27017', - 'database_name': 'test_db' + 'database_name': 'test_db', } @@ -147,21 +153,12 @@ def test_build_druid_config_defaults(): """Test that build_druid_config returns default values when no env vars are set""" config = build_druid_config() - assert config == { - 'host': 'localhost', - 'port': 8082 - } + assert config == {'host': 'localhost', 'port': 8082} def test_build_druid_config_with_env_vars(): """Test that build_druid_config uses env vars when set""" - with patch.dict(os.environ, { - 'DRUID_HOST': 'druid.example.com', - 'DRUID_PORT': '8083' - }): + with patch.dict(os.environ, {'DRUID_HOST': 'druid.example.com', 'DRUID_PORT': '8083'}): config = build_druid_config() - assert config == { - 'host': 'druid.example.com', - 'port': 8083 - } + assert config == {'host': 'druid.example.com', 'port': 8083} diff --git a/tests/workflow/sub_workflows/test_core_scouter.py b/tests/workflow/sub_workflows/test_core_scouter.py index 16ecff6..b60e19a 100644 --- a/tests/workflow/sub_workflows/test_core_scouter.py +++ b/tests/workflow/sub_workflows/test_core_scouter.py @@ -1,9 +1,11 @@ -from unittest.mock import AsyncMock, patch, call, ANY +from unittest.mock import ANY, AsyncMock, call, patch + import pytest -from scouter.workflow.sub_workflows.core_scouter import CoreScouter -from scouter.activities.activities import Activities from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ +from scouter.activities.activities import Activities +from scouter.workflow.sub_workflows.core_scouter import CoreScouter + @pytest.fixture def core_scouter(): @@ -14,7 +16,10 @@ def core_scouter(): @patch('scouter.workflow.sub_workflows.core_scouter.workflow', new_callable=AsyncMock) async def test_core_scouter_workflow_success(mock_workflow, core_scouter): mock_workflow.execute_local_activity_method.side_effect = [ - 'filtered_data', 'grouped_data', 'held_data'] + 'filtered_data', + 'grouped_data', + 'held_data', + ] await core_scouter.run( input_data={ 'metadata': { @@ -22,7 +27,7 @@ async def test_core_scouter_workflow_success(mock_workflow, core_scouter): 'model_id': 'test_model_id', 'model_name': 'test_model', 'schedule_name': 'test_schedule', - 'workflow_name': 'test_workflow' + 'workflow_name': 'test_workflow', } }, 'workflow_name': 'test_workflow', @@ -36,7 +41,7 @@ async def test_core_scouter_workflow_success(mock_workflow, core_scouter): 'table_name': 'test_table', 'retention_time': 3600, 'model_tags': {}, - 'debug_data_package': True + 'debug_data_package': True, } ) @@ -45,81 +50,90 @@ async def test_core_scouter_workflow_success(mock_workflow, core_scouter): 'model_id': 'test_model_id', 'model_name': 'test_model', 'schedule_name': 'test_schedule', - 'workflow_name': 'test_workflow' + 'workflow_name': 'test_workflow', } } - mock_workflow.execute_local_activity_method.assert_has_calls([ - call( - Activities.data_quality_gate, - { - **expected_metadata, - 'filters': {'test_filter': 'test_value'}, - 'data': 'test_data', - 'model_tags': {} - }, - retry_policy=ANY, - start_to_close_timeout=ANY - )]) - mock_workflow.execute_local_activity_method.assert_has_calls([ - call( - Activities.aggregate_data, - { - **expected_metadata, - 'data': 'filtered_data', - 'model_tags': {} - }, - retry_policy=ANY, - start_to_close_timeout=ANY - )]) - mock_workflow.execute_local_activity_method.assert_has_calls([ - call( - Activities.group_and_hold_data, - { - **expected_metadata, - 'workflow_name': 'test_workflow', - 'schedule_name': 'test_schedule', - 'data': 'grouped_data', - 'model_id': 'test_model_id', - 'retention_time': 3600, - 'model_tags': {} - }, - retry_policy=ANY, - start_to_close_timeout=ANY - )]) + mock_workflow.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.data_quality_gate, + { + **expected_metadata, + 'filters': {'test_filter': 'test_value'}, + 'data': 'test_data', + 'model_tags': {}, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + mock_workflow.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.aggregate_data, + {**expected_metadata, 'data': 'filtered_data', 'model_tags': {}}, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + mock_workflow.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.group_and_hold_data, + { + **expected_metadata, + 'workflow_name': 'test_workflow', + 'schedule_name': 'test_schedule', + 'data': 'grouped_data', + 'model_id': 'test_model_id', + 'retention_time': 3600, + 'model_tags': {}, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - mock_workflow.execute_activity_method.assert_has_calls([ - call( - Activities.export_data_to_postgres, - { - **expected_metadata, - 'schema': 'test_schema', - 'table_name': 'test_table', - 'data': 'held_data', - 'timestamp_conversion': { - 'column': 'timestamp', - 'format': DATETIME_FORMAT_WITH_TZ - } - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + mock_workflow.execute_activity_method.assert_has_calls( + [ + call( + Activities.export_data_to_postgres, + { + **expected_metadata, + 'schema': 'test_schema', + 'table_name': 'test_table', + 'data': 'held_data', + 'timestamp_conversion': { + 'column': 'timestamp', + 'format': DATETIME_FORMAT_WITH_TZ, + }, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) - mock_workflow.execute_activity_method.assert_has_calls([ - call( - Activities.store_data_package, - { - **expected_metadata, - 'workflow_name': 'test_workflow', - 'schedule_name': 'test_schedule', - 'held_data': 'held_data', - 'data': 'test_data' - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) + mock_workflow.execute_activity_method.assert_has_calls( + [ + call( + Activities.store_data_package, + { + **expected_metadata, + 'workflow_name': 'test_workflow', + 'schedule_name': 'test_schedule', + 'held_data': 'held_data', + 'data': 'test_data', + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) @pytest.mark.asyncio @@ -133,7 +147,7 @@ async def test_core_scouter_workflow_with_empty_data(mock_workflow, core_scouter 'model_id': 'test_model_id', 'model_name': 'test_model', 'schedule_name': 'test_schedule', - 'workflow_name': 'test_workflow' + 'workflow_name': 'test_workflow', } }, 'workflow_name': 'test_workflow', @@ -146,7 +160,7 @@ async def test_core_scouter_workflow_with_empty_data(mock_workflow, core_scouter 'schema': 'test_schema', 'table_name': 'test_table', 'retention_time': 3600, - 'model_tags': {} + 'model_tags': {}, } ) @@ -155,47 +169,52 @@ async def test_core_scouter_workflow_with_empty_data(mock_workflow, core_scouter 'model_id': 'test_model_id', 'model_name': 'test_model', 'schedule_name': 'test_schedule', - 'workflow_name': 'test_workflow' + 'workflow_name': 'test_workflow', } } - mock_workflow.execute_local_activity_method.assert_has_calls([ - call( - Activities.data_quality_gate, - { - **expected_metadata, - 'filters': {'test_filter': 'test_value'}, - 'data': 'test_data', - 'model_tags': {} - }, - retry_policy=ANY, - start_to_close_timeout=ANY - )]) - mock_workflow.execute_local_activity_method.assert_has_calls([ - call( - Activities.aggregate_data, - { - **expected_metadata, - 'data': {}, - 'model_tags': {} - }, - retry_policy=ANY, - start_to_close_timeout=ANY - )]) - mock_workflow.execute_local_activity_method.assert_has_calls([ - call( - Activities.group_and_hold_data, - { - **expected_metadata, - 'workflow_name': 'test_workflow', - 'schedule_name': 'test_schedule', - 'data': {}, - 'model_id': 'test_model_id', - 'retention_time': 3600, - 'model_tags': {} - }, - retry_policy=ANY, - start_to_close_timeout=ANY - )]) + mock_workflow.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.data_quality_gate, + { + **expected_metadata, + 'filters': {'test_filter': 'test_value'}, + 'data': 'test_data', + 'model_tags': {}, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + mock_workflow.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.aggregate_data, + {**expected_metadata, 'data': {}, 'model_tags': {}}, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + mock_workflow.execute_local_activity_method.assert_has_calls( + [ + call( + Activities.group_and_hold_data, + { + **expected_metadata, + 'workflow_name': 'test_workflow', + 'schedule_name': 'test_schedule', + 'data': {}, + 'model_id': 'test_model_id', + 'retention_time': 3600, + 'model_tags': {}, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) assert mock_workflow.execute_local_activity_method.call_count == 3 diff --git a/tests/workflow/test_fake_data.py b/tests/workflow/test_fake_data.py deleted file mode 100644 index f55e942..0000000 --- a/tests/workflow/test_fake_data.py +++ /dev/null @@ -1,29 +0,0 @@ -from unittest.mock import AsyncMock, patch, ANY -from pytest import fixture, mark -from scouter.workflow.fake_data import FakeData -from scouter.activities.faker import Faker - - -@fixture -def fake_data(): - return FakeData() - - -@mark.asyncio -@patch('scouter.workflow.fake_data.workflow', new_callable=AsyncMock) -async def test_fake_data_workflow(mock_workflow, fake_data): - mock_workflow.execute_activity_method.return_value = None - await fake_data.run( - { - 'topic': 'test_topic' - } - ) - - mock_workflow.execute_activity_method.assert_called_once_with( - Faker.generate_and_send_data, - { - 'topic': 'test_topic' - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) diff --git a/tests/workflow/test_scouter.py b/tests/workflow/test_scouter.py index 8d1b3a3..3e16666 100644 --- a/tests/workflow/test_scouter.py +++ b/tests/workflow/test_scouter.py @@ -1,7 +1,9 @@ -from unittest.mock import AsyncMock, patch, ANY, call +from unittest.mock import ANY, AsyncMock, call, patch + from pytest import fixture, mark -from scouter.workflow.scouter import Scouter + from scouter.activities.activities import Activities +from scouter.workflow.scouter import Scouter @fixture @@ -12,17 +14,16 @@ def scouter(): @mark.asyncio @patch('scouter.workflow.scouter.workflow', new_callable=AsyncMock) async def test_scouter_workflow(mock_workflow, scouter): - mock_workflow.execute_local_activity_method.side_effect = [ 'test_last_data_timestamp', - 'test_data' + 'test_data', ] await scouter.run( input_data={ 'topic': 'test_topic', 'schedule_name': 'test_schedule', 'model_name': 'test_model', - 'model_id': 'test_model_id' + 'model_id': 'test_model_id', } ) @@ -31,7 +32,7 @@ async def test_scouter_workflow(mock_workflow, scouter): 'model_id': 'test_model_id', 'model_name': 'test_model', 'schedule_name': 'test_schedule', - 'workflow_name': 'scouter' + 'workflow_name': 'scouter', } } @@ -39,15 +40,12 @@ async def test_scouter_workflow(mock_workflow, scouter): [ call( Activities.get_last_data_timestamp, - { - **expected_metadata, - 'workflow_name': 'scouter', - 'schedule_name': 'test_schedule' - }, + {**expected_metadata, 'workflow_name': 'scouter', 'schedule_name': 'test_schedule'}, retry_policy=ANY, - start_to_close_timeout=ANY + start_to_close_timeout=ANY, ) - ]) + ] + ) mock_workflow.execute_local_activity_method.assert_has_calls( [ @@ -55,13 +53,14 @@ async def test_scouter_workflow(mock_workflow, scouter): Activities.load_latest_data, { **expected_metadata, - 'collection_name': "raw_test_schedule", - 'last_data_timestamp': 'test_last_data_timestamp' + 'collection_name': 'raw_test_schedule', + 'last_data_timestamp': 'test_last_data_timestamp', }, retry_policy=ANY, - start_to_close_timeout=ANY + start_to_close_timeout=ANY, ) - ]) + ] + ) mock_workflow.execute_activity_method.assert_called_once_with( Activities.put_last_data_timestamp, @@ -69,10 +68,10 @@ async def test_scouter_workflow(mock_workflow, scouter): **expected_metadata, 'data': 'test_data', 'workflow_name': 'scouter', - 'schedule_name': 'test_schedule' + 'schedule_name': 'test_schedule', }, retry_policy=ANY, - start_to_close_timeout=ANY + start_to_close_timeout=ANY, ) mock_workflow.execute_child_workflow.assert_called_once_with( @@ -84,24 +83,21 @@ async def test_scouter_workflow(mock_workflow, scouter): 'workflow_name': 'scouter', 'schedule_name': 'test_schedule', 'model_name': 'test_model', - 'model_id': 'test_model_id' - } + 'model_id': 'test_model_id', + }, ) @mark.asyncio @patch('scouter.workflow.scouter.workflow', new_callable=AsyncMock) async def test_scouter_workflow_empty(mock_workflow, scouter): - mock_workflow.execute_local_activity_method.side_effect = [ - 'test_last_data_timestamp', - {} - ] + mock_workflow.execute_local_activity_method.side_effect = ['test_last_data_timestamp', {}] await scouter.run( input_data={ 'topic': 'test_topic', 'schedule_name': 'test_schedule', 'model_name': 'test_model', - 'model_id': 'test_model_id' + 'model_id': 'test_model_id', } ) @@ -110,7 +106,7 @@ async def test_scouter_workflow_empty(mock_workflow, scouter): 'model_id': 'test_model_id', 'model_name': 'test_model', 'schedule_name': 'test_schedule', - 'workflow_name': 'scouter' + 'workflow_name': 'scouter', } } @@ -120,13 +116,14 @@ async def test_scouter_workflow_empty(mock_workflow, scouter): Activities.load_latest_data, { **expected_metadata, - 'collection_name': "raw_test_schedule", - 'last_data_timestamp': 'test_last_data_timestamp' + 'collection_name': 'raw_test_schedule', + 'last_data_timestamp': 'test_last_data_timestamp', }, retry_policy=ANY, - start_to_close_timeout=ANY + start_to_close_timeout=ANY, ) - ]) + ] + ) mock_workflow.execute_activity_method.assert_not_called() mock_workflow.execute_child_workflow.assert_not_called() diff --git a/validate.sh b/validate.sh new file mode 100755 index 0000000..a80dca9 --- /dev/null +++ b/validate.sh @@ -0,0 +1,99 @@ +#!/bin/bash +# Model Manager Code Validation Script +# This script runs all code quality checks before committing or deploying + +set -e # Exit on any error + +# Colors for output +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +BLUE='\033[0;34m' +NC='\033[0m' # No Color + +echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}" +echo -e "${BLUE}║ Model Manager - Code Validation Suite ║${NC}" +echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}" +echo "" + +# Check if virtual environment is activated +if [[ -z "${VIRTUAL_ENV}" ]] && [[ -z "${CONDA_DEFAULT_ENV}" ]]; then + echo -e "${YELLOW}⚠️ Warning: No virtual environment detected${NC}" + echo -e "${YELLOW} Consider activating your venv/conda environment${NC}" + echo "" +fi + +# Function to run a validation step +run_step() { + local step_name=$1 + local step_command=$2 + + echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}" + echo -e "${BLUE}▶ ${step_name}${NC}" + echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}" + + if eval "$step_command"; then + echo -e "${GREEN}✅ ${step_name} - PASSED${NC}" + echo "" + return 0 + else + echo -e "${RED}❌ ${step_name} - FAILED${NC}" + echo "" + return 1 + fi +} + +# Track failures +FAILED_STEPS=() + +# Step 1: Code Formatting Check (Ruff) +if ! run_step "1. Code Formatting (Ruff)" "ruff format scouter/ tests/ && ruff format --check scouter/ tests/"; then + FAILED_STEPS+=("Code Formatting") +fi + +# Step 2: Linting (Ruff) +if ! run_step "2. Code Linting (Ruff)" "ruff check --fix scouter/ tests/"; then + FAILED_STEPS+=("Linting") +fi + +# Step 3: Type Checking (mypy) +if ! run_step "3. Type Checking (mypy)" "mypy scouter/"; then + FAILED_STEPS+=("Type Checking") +fi + +# Step 4: Security Analysis (Bandit) +if ! run_step "4. Security Analysis (Bandit)" "bandit -r scouter/ -ll -q"; then + FAILED_STEPS+=("Security Analysis") +fi + +# Step 5: Unit Tests (pytest) +if ! run_step "5. Unit Tests (pytest)" "pytest tests/ --cov=scouter --cov-report=term-missing --cov-report=xml --cov-report=html --cov-fail-under=80 -q"; then + FAILED_STEPS+=("Unit Tests") +fi + +# Summary +echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}" +echo -e "${BLUE}║ Validation Summary ║${NC}" +echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}" +echo "" + +if [ ${#FAILED_STEPS[@]} -eq 0 ]; then + echo -e "${GREEN}✅ All validation checks passed!${NC}" + echo -e "${GREEN} Your code is ready for commit/deployment.${NC}" + echo "" + exit 0 +else + echo -e "${RED}❌ Validation failed for the following steps:${NC}" + for step in "${FAILED_STEPS[@]}"; do + echo -e "${RED} • ${step}${NC}" + done + echo "" + echo -e "${YELLOW}💡 Tips:${NC}" + echo -e "${YELLOW} • Run 'ruff format scouter/ tests/' to auto-fix formatting${NC}" + echo -e "${YELLOW} • Run 'ruff check --fix scouter/ tests/' to auto-fix linting issues${NC}" + echo -e "${YELLOW} • Review mypy errors and add type hints where needed${NC}" + echo -e "${YELLOW} • Check bandit warnings for security issues${NC}" + echo -e "${YELLOW} • Fix failing tests or improve test coverage${NC}" + echo "" + exit 1 +fi \ No newline at end of file