From eeaf6a2a5b121fd8c39132a2c967369e1a3d080e Mon Sep 17 00:00:00 2001 From: Tal Borenstein Date: Mon, 14 Oct 2024 19:22:50 +0300 Subject: [PATCH] feat: test --- keep/contextmanager/contextmanager.py | 41 +++++++++--------- keep/providers/base/base_provider.py | 3 +- tests/test_workflow_execution.py | 62 ++++++++++++++++++++++++--- 3 files changed, 77 insertions(+), 29 deletions(-) diff --git a/keep/contextmanager/contextmanager.py b/keep/contextmanager/contextmanager.py index eeca97560..ce55e2cda 100644 --- a/keep/contextmanager/contextmanager.py +++ b/keep/contextmanager/contextmanager.py @@ -10,13 +10,13 @@ class ContextManager: def __init__(self, tenant_id, workflow_id=None, workflow_execution_id=None): - + self.logger = logging.getLogger(__name__) + self.logger_adapter = WorkflowLoggerAdapter( + self.logger, self, tenant_id, workflow_id, workflow_execution_id + ) self.workflow_id = workflow_id self.workflow_execution_id = workflow_execution_id self.tenant_id = tenant_id - self.logger = None - self.logger_adapter = None - self.set_logger(logging.getLogger(__name__)) self.steps_context = {} self.steps_context_size = 0 self.providers_context = {} @@ -54,6 +54,7 @@ def __init__(self, tenant_id, workflow_id=None, workflow_execution_id=None): self.dependencies = set() self.workflow_execution_id = None self._api_key = None + self.__loggers = [] @property def api_key(self): @@ -74,23 +75,23 @@ def api_key(self): def set_execution_context(self, workflow_execution_id): self.workflow_execution_id = workflow_execution_id self.logger_adapter.workflow_execution_id = workflow_execution_id - - def set_logger_by_name(self, logger_name): - logger = logging.getLogger(logger_name) - log_level_string = os.environ.get("KEEP_PROVIDER_{}_LOG_LEVEL".format(logger_name.upper()), None) - if log_level_string: - log_level = logging.getLevelName(log_level_string) - logger.setLevel(log_level) - self.set_logger(logger) - - def set_logger(self, logger): - self.logger = logger - self.logger_adapter = WorkflowLoggerAdapter( - self.logger, self, self.tenant_id, self.workflow_id, self.workflow_execution_id + for logger in self.__loggers: + logger.workflow_execution_id = workflow_execution_id + + def get_logger(self, name=None): + if not name: + return self.logger_adapter + + logger = logging.getLogger(name) + logger_adapter = WorkflowLoggerAdapter( + logger, + self, + self.tenant_id, + self.workflow_id, + self.workflow_execution_id, ) - - def get_logger(self): - return self.logger_adapter + self.__loggers.append(logger_adapter) + return logger_adapter def set_event_context(self, event): self.event_context = event diff --git a/keep/providers/base/base_provider.py b/keep/providers/base/base_provider.py index d9e21fc7e..844524551 100644 --- a/keep/providers/base/base_provider.py +++ b/keep/providers/base/base_provider.py @@ -66,8 +66,7 @@ def __init__( self.webhook_markdown = webhook_markdown self.provider_description = provider_description self.context_manager = context_manager - context_manager.set_logger_by_name("provider_{}".format(self.provider_id)) - self.logger = context_manager.get_logger() + self.logger = context_manager.get_logger(self.__class__.__name__) self.validate_config() self.logger.debug( "Base provider initalized", extra={"provider": self.__class__.__name__} diff --git a/tests/test_workflow_execution.py b/tests/test_workflow_execution.py index 17910e42a..b55ad118b 100644 --- a/tests/test_workflow_execution.py +++ b/tests/test_workflow_execution.py @@ -44,6 +44,33 @@ """ +workflow_definition_with_two_providers = """workflow: +id: susu-and-sons +description: Just to test the logs of 2 providers +triggers: +- type: alert + filters: + - key: name + value: "server-is-hamburger" +steps: +- name: keep_step + provider: + type: keep + with: + filters: + - key: status + value: open +actions: +- name: console_action + provider: + type: console + with: + message: | + "Tier 1 Alert: {{ alert.name }} - {{ alert.description }} + Alert details: {{ alert }}" +""" + + @pytest.fixture(scope="module") def workflow_manager(): """ @@ -77,6 +104,25 @@ def setup_workflow(db_session): db_session.commit() +@pytest.fixture +def setup_workflow_with_two_providers(db_session): + """ + Fixture to set up a workflow in the database before each test. + It creates a Workflow object with the predefined workflow definition and adds it to the database. + """ + workflow = Workflow( + id="susu-and-sons", + name="susu-and-sons", + tenant_id=SINGLE_TENANT_UUID, + description="some stuff for unit testing", + created_by="tal@keephq.dev", + interval=0, + workflow_raw=workflow_definition_with_two_providers, + ) + db_session.add(workflow) + db_session.commit() + + @pytest.mark.parametrize( "test_app, test_case, alert_statuses, expected_tier, db_session", [ @@ -807,7 +853,7 @@ def test_workflow_execution_logs( db_session, test_app, create_alert, - setup_workflow, + setup_workflow_with_two_providers, workflow_manager, test_case, alert_statuses, @@ -828,7 +874,7 @@ def test_workflow_execution_logs( current_alert = AlertDto( id="grafana-1", source=["grafana"], - name="server-is-down", + name="server-is-hamburger", status=AlertStatus.FIRING, severity="critical", fingerprint="fp1", @@ -843,7 +889,7 @@ def test_workflow_execution_logs( status = None while workflow_execution is None and count < 30 and status != "success": workflow_execution = get_last_workflow_execution_by_workflow_id( - SINGLE_TENANT_UUID, "alert-time-check" + SINGLE_TENANT_UUID, "susu-and-sons" ) if workflow_execution is not None: status = workflow_execution.status @@ -854,8 +900,10 @@ def test_workflow_execution_logs( assert workflow_execution is not None assert workflow_execution.status == "success" - logs = db_session.query(WorkflowExecutionLog).filter( - WorkflowExecutionLog.workflow_execution_id == workflow_execution.id - ).all() + logs = ( + db_session.query(WorkflowExecutionLog) + .filter(WorkflowExecutionLog.workflow_execution_id == workflow_execution.id) + .all() + ) - assert len(logs) == 4 \ No newline at end of file + assert len(logs) == 15