diff --git a/.gitignore b/.gitignore index c503b7f..c7d6918 100644 --- a/.gitignore +++ b/.gitignore @@ -39,6 +39,7 @@ __pycache__/ # Ignorar coverage htmlcov/ .coverage +coverage.xml # git keys git_key* diff --git a/orchestrator/activities/formatters.py b/orchestrator/activities/formatters.py index d32b4c0..0584ce6 100644 --- a/orchestrator/activities/formatters.py +++ b/orchestrator/activities/formatters.py @@ -1,3 +1,5 @@ +from collections.abc import Hashable + from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): @@ -87,7 +89,10 @@ class Formatters(BaseActivity): pipelines = input_data['pipelines'] - schedule_config = {self.scouter_namespace: {}, self.laborious_namespace: {}} + schedule_config: dict[str, dict[str, Any]] = { + self.scouter_namespace: {}, + self.laborious_namespace: {}, + } for pipeline in pipelines: if pipeline['workflow_type'] == 'scouter': @@ -155,7 +160,7 @@ class Formatters(BaseActivity): number_of_slots = len(active_ingestors) if active_ingestors else 1 tags_per_slot = ceil(number_of_tags / number_of_slots) - slot_config = {} + slot_config: dict[str, Any] = {} last_index = 0 for i in range(1, number_of_slots): @@ -212,7 +217,7 @@ class Formatters(BaseActivity): schedule_config = input_data['schedule_config'] - config = {} + config: dict[str, dict[str, str]] = {} for schedule in schedule_config: namespace = schedule['namespace'] schedule_name = schedule['schedule_name'] @@ -290,9 +295,18 @@ class Formatters(BaseActivity): current_schedule_config = input_data['current_schedule_config'] schedule_config = input_data['schedule_config'] - to_update = {self.scouter_namespace: {}, self.laborious_namespace: {}} - to_create = {self.scouter_namespace: {}, self.laborious_namespace: {}} - to_delete = {self.scouter_namespace: [], self.laborious_namespace: []} + to_update: dict[str, dict[str, Any]] = { + self.scouter_namespace: {}, + self.laborious_namespace: {}, + } + to_create: dict[str, dict[str, Any]] = { + self.scouter_namespace: {}, + self.laborious_namespace: {}, + } + to_delete: dict[str, list[str]] = { + self.scouter_namespace: [], + self.laborious_namespace: [], + } for namespace, schedules in schedule_config.items(): current_schedules = current_schedule_config.get(namespace, {}) @@ -352,7 +366,11 @@ class Formatters(BaseActivity): return output def send_success_report( - self, metadata: dict[str, Any], message: str, notification_id: str, attachment: str = None + self, + metadata: dict[str, Any], + message: str, + notification_id: str, + attachment: Any | None = None, ) -> None: """ Sends a success notification report. @@ -393,7 +411,9 @@ class Formatters(BaseActivity): attachment_content=attachment, ) - def parse_report_schedule(self, input_data: dict[str, Any]) -> tuple[list[str], dict[str, Any]]: + def parse_report_schedule( + self, input_data: list[dict[str, Any]] + ) -> tuple[list[str], dict[str, Any]]: """ Parses the report schedule data to extract success and error information. @@ -424,7 +444,7 @@ class Formatters(BaseActivity): return success_keys, error_keys - def parse_report(self, input_data: dict[str, Any]) -> tuple[list[str], list[str]]: + def parse_report(self, input_data: dict[str, dict[str, Any]]) -> tuple[list[str], list[str]]: """ Parses the report data to extract success and error keys. @@ -587,7 +607,7 @@ class Formatters(BaseActivity): ) @activity.defn(name='format_log_report') - async def format_log_report(self, input_data: dict[str, Any]) -> dict[str, Any]: + async def format_log_report(self, input_data: dict[str, Any]) -> dict[Hashable, Any]: """ Formats the receiver_groups status to a dataframe to be stored in the database. @@ -636,7 +656,9 @@ class Formatters(BaseActivity): if group_name not in data[key]['groups']: data[key]['groups'].append(group_name) - return DataFrame(list(data.values())).to_dict() + data_values: DataFrame = DataFrame(list(data.values())) + + return data_values.to_dict() @activity.defn(name='filter_notification_reports') async def filter_notification_reports(self, input_data: dict[str, Any]) -> dict[str, Any]: diff --git a/orchestrator/activities/mongo_db.py b/orchestrator/activities/mongo_db.py index 1dfba8a..8f1969f 100644 --- a/orchestrator/activities/mongo_db.py +++ b/orchestrator/activities/mongo_db.py @@ -70,7 +70,9 @@ 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] diff --git a/orchestrator/activities/temporal_manager.py b/orchestrator/activities/temporal_manager.py index 6e7f75e..1de072f 100644 --- a/orchestrator/activities/temporal_manager.py +++ b/orchestrator/activities/temporal_manager.py @@ -52,7 +52,7 @@ class TemporalManager(BaseActivity): self.temporal_host = host self.scouter_namespace = scouter_namespace self.laborious_namespace = laborious_namespace - self.temporal_clients = {} + self.temporal_clients: dict[str, Client] = {} self.model_id_id_key = SearchAttributeKey.for_keyword('model_id') self.model_name_id_key = SearchAttributeKey.for_keyword('model_name') @@ -84,7 +84,7 @@ class TemporalManager(BaseActivity): } @activity.defn(name='normalize_schedules') - async def normalize_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]: + async def normalize_schedules(self, input_data: dict[str, Any]): """ Normalize schedules. Removes schedules with no update time in mongo db collection "orchestrated_schedules". input_data: diff --git a/orchestrator/utils/email_builder.py b/orchestrator/utils/email_builder.py index 8c22203..2a815e5 100644 --- a/orchestrator/utils/email_builder.py +++ b/orchestrator/utils/email_builder.py @@ -1,3 +1,5 @@ +from typing import Any + from jinja2 import Template from sientia_do.observability.logger import Logger @@ -75,12 +77,12 @@ class EmailBuilder: else '', } - def build_email(self, report_data: list[dict], mail_type: str) -> str: + def build_email(self, report_data: list[dict[str, Any]], mail_type: str) -> str: """ Builds the email HTML by organizing report data by notification level and model. Args: - report_data (list[dict]): List of notification reports, each containing: + report_data (List[Dict[str, Any]]): List of notification reports, each containing: - level (str): Notification level (ERROR, WARNING, INFO) - model_name (str): Name of the model - Additional notification details @@ -88,7 +90,7 @@ class EmailBuilder: Returns: str: Complete HTML email content ready for sending. """ - general_events = {} + general_events: dict[str, dict[str, Any]] = {} for report in report_data: level = report['level'] @@ -100,13 +102,17 @@ class EmailBuilder: 'models': {}, } - if model_name not in general_events[level]['models']: - general_events[level]['models'][model_name] = { + # Type assertion to help the type checker understand the structure + level_data = general_events[level] + models_dict = level_data['models'] + + if model_name not in models_dict: + models_dict[model_name] = { 'model_name': model_name, 'events': [], } - general_events[level]['models'][model_name]['events'].append(report) + models_dict[model_name]['events'].append(report) for _type, content in general_events.items(): content['models'] = list(content['models'].values())