use crate::dispatch::error::DispatchError;
use scouter_types::{
AlertDispatchConfig, AlertDispatchType, DispatchAlertDescription, DispatchDriftConfig,
DriftArgs, OpsGenieDispatchConfig, SlackDispatchConfig,
};
use serde_json::{json, Value};
use std::result::Result;
use std::{collections::HashMap, env};
use tracing::error;
trait DispatchHelpers {
fn construct_alert_description<T: DispatchAlertDescription>(
&self,
feature_alerts: &T,
) -> String;
}
pub trait Dispatch {
fn process_alerts<T: DispatchAlertDescription + std::marker::Sync>(
&self,
feature_alerts: &T,
) -> impl std::future::Future<Output = Result<(), DispatchError>> + Send;
}
pub trait HttpAlertWrapper {
fn api_url(&self) -> &str;
fn header_auth_value(&self) -> &str;
fn construct_alert_body(&self, alert_description: &str) -> Value;
}
#[derive(Debug)]
pub struct OpsGenieAlerter {
header_auth_value: String,
api_url: String,
team_name: String,
priority: String,
name: String,
space: String,
version: String,
}
impl OpsGenieAlerter {
pub fn new(
name: &str,
space: &str,
version: &str,
dispatch_config: &OpsGenieDispatchConfig,
) -> Result<Self, DispatchError> {
let api_key = env::var("OPSGENIE_API_KEY")
.map_err(|_| DispatchError::OpsGenieError("OPSGENIE_API_KEY is not set".to_string()))?;
let api_url = env::var("OPSGENIE_API_URL")
.map_err(|_| DispatchError::OpsGenieError("OPSGENIE_API_URL is not set".to_string()))?;
let team_name = dispatch_config.team.clone();
let priority = dispatch_config.priority.clone();
Ok(Self {
header_auth_value: format!("GenieKey {api_key}"),
api_url,
team_name,
name: name.to_string(),
space: space.to_string(),
version: version.to_string(),
priority,
})
}
}
impl HttpAlertWrapper for OpsGenieAlerter {
fn api_url(&self) -> &str {
&self.api_url
}
fn header_auth_value(&self) -> &str {
&self.header_auth_value
}
fn construct_alert_body(&self, alert_description: &str) -> Value {
let mut mapping: HashMap<&str, Value> = HashMap::new();
mapping.insert(
"message",
format!(
"Model drift detected for {}/{}/{}",
self.space, self.name, self.version
)
.into(),
);
mapping.insert("description", alert_description.to_string().into());
mapping.insert(
"responders",
json!([{"name": self.team_name, "type": "team"}]),
);
mapping.insert(
"visibleTo",
json!([{"name": self.team_name, "type": "team"}]),
);
mapping.insert("tags", json!(["Model Drift", "Scouter"]));
mapping.insert("priority", self.priority.clone().into());
json!(mapping)
}
}
impl DispatchHelpers for OpsGenieAlerter {
fn construct_alert_description<T: DispatchAlertDescription>(
&self,
feature_alerts: &T,
) -> String {
feature_alerts.create_alert_description(AlertDispatchType::OpsGenie)
}
}
#[derive(Debug)]
pub struct SlackAlerter {
header_auth_value: String,
api_url: String,
name: String,
space: String,
version: String,
channel: String,
}
impl SlackAlerter {
pub fn new(
name: &str,
space: &str,
version: &str,
dispatch_config: &SlackDispatchConfig,
) -> Result<Self, DispatchError> {
let app_token = env::var("SLACK_APP_TOKEN")
.map_err(|_| DispatchError::SlackError("SLACK_APP_TOKEN not set".to_string()))?;
let api_url = env::var("SLACK_API_URL")
.map_err(|_| DispatchError::SlackError("SLACK_API_URL not set".to_string()))?;
let slack_channel = dispatch_config.channel.clone();
Ok(Self {
header_auth_value: format!("Bearer {app_token}"),
api_url: format!("{api_url}/chat.postMessage"),
name: name.to_string(),
space: space.to_string(),
version: version.to_string(),
channel: slack_channel,
})
}
}
impl HttpAlertWrapper for SlackAlerter {
fn api_url(&self) -> &str {
&self.api_url
}
fn header_auth_value(&self) -> &str {
&self.header_auth_value
}
fn construct_alert_body(&self, alert_description: &str) -> Value {
json!({
"channel": self.channel,
"blocks": [
{
"type": "header",
"text": {
"type": "plain_text",
"text": ":rotating_light: Drift Detected :rotating_light:",
"emoji": true
}
},
{
"type": "section",
"text": {
"type": "mrkdwn",
"text": format!("*Name*: {} *Space*: {} *Version*: {}", self.name, self.space, self.version),
}
},
{
"type": "section",
"text": {
"type": "mrkdwn",
"text": alert_description
},
}
]
})
}
}
impl DispatchHelpers for SlackAlerter {
fn construct_alert_description<T: DispatchAlertDescription>(
&self,
feature_alerts: &T,
) -> String {
feature_alerts.create_alert_description(AlertDispatchType::Slack)
}
}
#[derive(Debug)]
pub struct HttpAlertDispatcher<T: HttpAlertWrapper> {
http_client: reqwest::Client,
alerter: T,
}
impl<T: HttpAlertWrapper> HttpAlertDispatcher<T> {
pub fn new(alerter: T) -> Self {
Self {
http_client: reqwest::Client::new(),
alerter,
}
}
async fn send_alerts(&self, body: Value) -> Result<(), DispatchError> {
let response = self
.http_client
.post(self.alerter.api_url())
.header("Authorization", self.alerter.header_auth_value())
.json(&body)
.send()
.await
.map_err(|e| DispatchError::HttpError(e.to_string()))?;
if response.status().is_success() {
Ok(())
} else {
let text = response
.text()
.await
.unwrap_or("Failed to parse response".to_string());
error!("Failed to send alert: {}. Continuing", text);
Ok(())
}
}
}
impl<T: HttpAlertWrapper + DispatchHelpers + std::marker::Sync> Dispatch
for HttpAlertDispatcher<T>
{
async fn process_alerts<J: DispatchAlertDescription>(
&self,
feature_alerts: &J,
) -> Result<(), DispatchError> {
let alert_description = self.alerter.construct_alert_description(feature_alerts);
let alert_body = self.alerter.construct_alert_body(&alert_description);
self.send_alerts(alert_body)
.await
.map_err(|e| DispatchError::HttpError(format!("Failed to send alerts: {e}")))?;
Ok(())
}
}
#[derive(Debug)]
pub struct ConsoleAlertDispatcher {
name: String,
space: String,
version: String,
}
impl ConsoleAlertDispatcher {
pub fn new(name: &str, space: &str, version: &str) -> Self {
Self {
name: name.to_string(),
space: space.to_string(),
version: version.to_string(),
}
}
}
impl Dispatch for ConsoleAlertDispatcher {
async fn process_alerts<T: DispatchAlertDescription>(
&self,
feature_alerts: &T,
) -> Result<(), DispatchError> {
let alert_description = self.construct_alert_description(feature_alerts);
if !alert_description.is_empty() {
let msg1 = "Drift detected for";
let msg2 = format!("{}/{}/{}!", self.space, self.name, self.version);
let mut body = format!("\n{msg1} {msg2} \n");
body.push_str(&alert_description);
println!("{body}");
}
Ok(())
}
}
impl DispatchHelpers for ConsoleAlertDispatcher {
fn construct_alert_description<T: DispatchAlertDescription>(
&self,
feature_alerts: &T,
) -> String {
feature_alerts.create_alert_description(AlertDispatchType::Console)
}
}
#[derive(Debug)]
pub enum AlertDispatcher {
Console(ConsoleAlertDispatcher),
OpsGenie(HttpAlertDispatcher<OpsGenieAlerter>),
Slack(HttpAlertDispatcher<SlackAlerter>),
}
impl AlertDispatcher {
pub async fn process_alerts<T: DispatchAlertDescription + std::marker::Sync>(
&self,
feature_alerts: &T,
) -> Result<(), DispatchError> {
match self {
AlertDispatcher::Console(dispatcher) => dispatcher
.process_alerts(feature_alerts)
.await
.map_err(|e| DispatchError::AlertProcessError(e.to_string())),
AlertDispatcher::OpsGenie(dispatcher) => dispatcher
.process_alerts(feature_alerts)
.await
.map_err(|e| DispatchError::AlertProcessError(e.to_string())),
AlertDispatcher::Slack(dispatcher) => dispatcher
.process_alerts(feature_alerts)
.await
.map_err(|e| DispatchError::AlertProcessError(e.to_string())),
}
}
pub fn new<T: DispatchDriftConfig>(config: &T) -> Result<Self, DispatchError> {
let args: DriftArgs = config.get_drift_args();
let result = match args.dispatch_config {
AlertDispatchConfig::Slack(config) => {
SlackAlerter::new(&args.name, &args.space, &args.version, &config)
.map(|alerter| AlertDispatcher::Slack(HttpAlertDispatcher::new(alerter)))
}
AlertDispatchConfig::OpsGenie(config) => {
OpsGenieAlerter::new(&args.name, &args.space, &args.version, &config)
.map(|alerter| AlertDispatcher::OpsGenie(HttpAlertDispatcher::new(alerter)))
}
AlertDispatchConfig::Console(_) => Ok(AlertDispatcher::Console(
ConsoleAlertDispatcher::new(&args.name, &args.space, &args.version),
)),
};
match result {
Ok(dispatcher) => Ok(dispatcher),
Err(e) => {
error!("Failed to create Alerter: {:?}", e);
Ok(AlertDispatcher::Console(ConsoleAlertDispatcher::new(
&args.name,
&args.space,
&args.version,
)))
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use scouter_types::spc::{
AlertZone, SpcAlert, SpcAlertConfig, SpcAlertType, SpcDriftConfig, SpcFeatureAlert,
SpcFeatureAlerts,
};
use std::collections::HashMap;
use std::env;
fn test_features_map() -> HashMap<String, SpcFeatureAlert> {
let mut features: HashMap<String, SpcFeatureAlert> = HashMap::new();
features.insert(
"test_feature_1".to_string(),
SpcFeatureAlert {
feature: "test_feature_1".to_string(),
alerts: vec![SpcAlert {
zone: AlertZone::Zone4,
kind: SpcAlertType::OutOfBounds,
}]
.into_iter()
.collect(),
},
);
features.insert(
"test_feature_2".to_string(),
SpcFeatureAlert {
feature: "test_feature_2".to_string(),
alerts: vec![SpcAlert {
zone: AlertZone::Zone1,
kind: SpcAlertType::Consecutive,
}]
.into_iter()
.collect(),
},
);
features
}
#[test]
fn test_construct_opsgenie_alert_description() {
unsafe {
env::set_var("OPSGENIE_API_URL", "api_url");
env::set_var("OPSGENIE_API_KEY", "api_key");
}
let features = test_features_map();
let alerter = OpsGenieAlerter::new(
"name",
"space",
"1.0.0",
&OpsGenieDispatchConfig {
team: "test-team".to_string(),
priority: "P5".to_string(),
},
)
.unwrap();
let alert_description = alerter.construct_alert_description(&SpcFeatureAlerts {
features,
has_alerts: true,
});
let expected_alert_description = "Drift has been detected for the following features:\n test_feature_2: \n Kind: Consecutive\n Zone: Zone 1\n test_feature_1: \n Kind: Out of bounds\n Zone: Zone 4\n".to_string();
assert_eq!(&alert_description.len(), &expected_alert_description.len());
unsafe {
env::remove_var("OPSGENIE_API_URL");
env::remove_var("OPSGENIE_API_KEY");
}
}
#[test]
fn test_construct_opsgenie_alert_description_empty() {
unsafe {
env::set_var("OPSGENIE_API_URL", "api_url");
env::set_var("OPSGENIE_API_KEY", "api_key");
}
let features: HashMap<String, SpcFeatureAlert> = HashMap::new();
let alerter = OpsGenieAlerter::new(
"name",
"space",
"1.0.0",
&OpsGenieDispatchConfig {
team: "test-team".to_string(),
priority: "P5".to_string(),
},
)
.unwrap();
let alert_description = alerter.construct_alert_description(&SpcFeatureAlerts {
features,
has_alerts: true,
});
let expected_alert_description = "".to_string();
assert_eq!(alert_description, expected_alert_description);
unsafe {
env::remove_var("OPSGENIE_API_URL");
env::remove_var("OPSGENIE_API_KEY");
}
}
#[tokio::test]
async fn test_construct_opsgenie_alert_body() {
let download_server = mockito::Server::new_async().await;
let url = download_server.url();
unsafe {
env::set_var("OPSGENIE_API_URL", url);
env::set_var("OPSGENIE_API_KEY", "api_key");
}
let ops_genie_team = "test-team";
let expected_alert_body = json!(
{
"message": "Model drift detected for test_repo/test_ml_model/1.0.0",
"description": "Features have drifted",
"responders":[
{"name":ops_genie_team, "type":"team"}
],
"visibleTo":[
{"name":ops_genie_team, "type":"team"}
],
"tags": ["Model Drift", "Scouter"],
"priority": "P1"
}
);
let alerter = OpsGenieAlerter::new(
"test_ml_model",
"test_repo",
"1.0.0",
&OpsGenieDispatchConfig {
team: ops_genie_team.to_string(),
priority: "P1".to_string(),
},
)
.unwrap();
let alert_body = alerter.construct_alert_body("Features have drifted");
assert_eq!(alert_body, expected_alert_body);
unsafe {
env::remove_var("OPSGENIE_API_URL");
env::remove_var("OPSGENIE_API_KEY");
}
}
#[tokio::test]
async fn test_send_opsgenie_alerts() {
let mut download_server = mockito::Server::new_async().await;
let url = format!("{}/alerts", download_server.url());
unsafe {
env::set_var("OPSGENIE_API_URL", url);
env::set_var("OPSGENIE_API_KEY", "api_key");
}
let mock_get_path = download_server
.mock("Post", "/alerts")
.with_status(201)
.create();
let features = test_features_map();
let dispatcher = AlertDispatcher::OpsGenie(HttpAlertDispatcher::new(
OpsGenieAlerter::new(
"name",
"space",
"1.0.0",
&OpsGenieDispatchConfig {
team: "test-team".to_string(),
priority: "P5".to_string(),
},
)
.unwrap(),
));
let _ = dispatcher
.process_alerts(&SpcFeatureAlerts {
features,
has_alerts: true,
})
.await;
mock_get_path.assert();
unsafe {
env::remove_var("OPSGENIE_API_URL");
env::remove_var("OPSGENIE_API_KEY");
}
}
#[tokio::test]
async fn test_send_console_alerts() {
let features = test_features_map();
let dispatcher =
AlertDispatcher::Console(ConsoleAlertDispatcher::new("name", "space", "1.0.0"));
let result = dispatcher
.process_alerts(&SpcFeatureAlerts {
features,
has_alerts: true,
})
.await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_send_slack_alerts() {
let mut download_server = mockito::Server::new_async().await;
let url = download_server.url();
unsafe {
env::set_var("SLACK_API_URL", url);
env::set_var("SLACK_APP_TOKEN", "bot_token");
}
let mock_get_path = download_server
.mock("Post", "/chat.postMessage")
.with_status(201)
.create();
let features = test_features_map();
let dispatcher = AlertDispatcher::Slack(HttpAlertDispatcher::new(
SlackAlerter::new(
"name",
"space",
"1.0.0",
&SlackDispatchConfig {
channel: "test-channel".to_string(),
},
)
.unwrap(),
));
let _ = dispatcher
.process_alerts(&SpcFeatureAlerts {
features,
has_alerts: true,
})
.await;
mock_get_path.assert();
unsafe {
env::remove_var("SLACK_API_URL");
env::remove_var("SLACK_APP_TOKEN");
}
}
#[tokio::test]
async fn test_construct_slack_alert_body() {
let download_server = mockito::Server::new_async().await;
let url = download_server.url();
let slack_channel = "test_channel";
unsafe {
env::set_var("SLACK_API_URL", url);
env::set_var("SLACK_APP_TOKEN", "bot_token");
}
let expected_alert_body = json!({
"channel": slack_channel,
"blocks": [
{
"type": "header",
"text": {
"type": "plain_text",
"text": ":rotating_light: Drift Detected :rotating_light:",
"emoji": true
}
},
{
"type": "section",
"text": {
"type": "mrkdwn",
"text": "*Name*: name *Space*: space *Version*: 1.0.0",
}
},
{
"type": "section",
"text": {
"type": "mrkdwn",
"text": "*Features have drifted*"
},
}
]
});
let alerter = SlackAlerter::new(
"name",
"space",
"1.0.0",
&SlackDispatchConfig {
channel: slack_channel.to_string(),
},
)
.unwrap();
let alert_body = alerter.construct_alert_body("*Features have drifted*");
assert_eq!(alert_body, expected_alert_body);
unsafe {
env::remove_var("SLACK_API_URL");
env::remove_var("SLACK_APP_TOKEN");
}
}
#[test]
fn test_console_dispatcher_returned_when_env_vars_not_set_opsgenie() {
unsafe {
env::remove_var("OPSGENIE_API_KEY");
}
let alert_config = SpcAlertConfig {
dispatch_config: AlertDispatchConfig::OpsGenie(OpsGenieDispatchConfig {
team: "test-team".to_string(),
priority: "P5".to_string(),
}),
..Default::default()
};
let config = SpcDriftConfig::new(
"name",
"space",
"1.0.0",
None,
None,
Some(alert_config),
None,
)
.unwrap();
let dispatcher = AlertDispatcher::new(&config).unwrap();
assert!(
matches!(dispatcher, AlertDispatcher::Console(_)),
"Expected Console Dispatcher"
);
}
#[test]
fn test_console_dispatcher_returned_when_env_vars_not_set_slack() {
unsafe {
env::remove_var("SLACK_API_URL");
env::remove_var("SLACK_APP_TOKEN");
}
let alert_config = SpcAlertConfig {
dispatch_config: AlertDispatchConfig::Slack(SlackDispatchConfig {
channel: "test-channel".to_string(),
}),
..Default::default()
};
let config = SpcDriftConfig::new(
"name",
"space",
"1.0.0",
None,
None,
Some(alert_config),
None,
)
.unwrap();
let dispatcher = AlertDispatcher::new(&config).unwrap();
assert!(
matches!(dispatcher, AlertDispatcher::Console(_)),
"Expected Console Dispatcher"
);
}
#[test]
fn test_slack_dispatcher_returned_when_env_vars_set() {
unsafe {
env::set_var("SLACK_API_URL", "url");
env::set_var("SLACK_APP_TOKEN", "bot_token");
}
let alert_config = SpcAlertConfig {
dispatch_config: AlertDispatchConfig::Slack(SlackDispatchConfig {
channel: "test-channel".to_string(),
}),
..Default::default()
};
let config = SpcDriftConfig::new(
"name",
"space",
"1.0.0",
None,
None,
Some(alert_config),
None,
)
.unwrap();
let dispatcher = AlertDispatcher::new(&config).unwrap();
assert!(
matches!(dispatcher, AlertDispatcher::Slack(_)),
"Expected Slack Dispatcher"
);
unsafe {
env::remove_var("SLACK_API_URL");
env::remove_var("SLACK_APP_TOKEN");
}
}
}