patch: adds host arg (#15)

Author: jtyoung84

src/aind_airflow_jobs/alert_handler.py

@@ -82,6 +82,7 @@ def send_log_message(
     job_info: Dict[str, Any],
     message: Optional[str] = None,
     log_level: str = "INFO",
+    loki_host: Optional[str] = None,
 ) -> None:
     """
     Sends a log message given the job information.
@@ -90,6 +91,8 @@ def send_log_message(
     job_info : Dict[str, Any]
     message : str | None
     log_level : str
+    loki_host : str | None
+      Optional override of loki host
 
     Returns
     -------
@@ -100,7 +103,7 @@ def send_log_message(
     job_name = job_info["job_name"]
     run_id = job_info["run_id"]
     task_id = job_info["task_id"]
-    handler = LokiLoggerHook().get_conn()
+    handler = LokiLoggerHook(host=loki_host).get_conn()
     logging_message = (
         f"{job_name} on {run_id} and {task_id}"
         if message is None

tests/test_alert_handler.py

@@ -85,7 +85,7 @@ def test_send_log_message(
         mock_logger.return_value.log = mock_log
         send_log_message(job_info=job_info, log_level="DEBUG", message="Hello")
 
-        mock_loki.assert_has_calls([call(), call().get_conn()])
+        mock_loki.assert_has_calls([call(host=None), call().get_conn()])
         mock_logger.assert_has_calls([call("aind_airflow_jobs.alert_handler")])
         mock_log.assert_called_once_with(
             10,