def test_stackdriver_disable_notification_channel(
self, mock_channel_client, mock_get_creds_and_project_id):
hook = stackdriver.StackdriverHook()
notification_channel_enabled = NotificationChannel(
**TEST_NOTIFICATION_CHANNEL_1)
notification_channel_disabled = NotificationChannel(
**TEST_NOTIFICATION_CHANNEL_2)
mock_channel_client.return_value.list_notification_channels.return_value = [
notification_channel_enabled,
notification_channel_disabled,
]
hook.disable_notification_channels(
filter_=TEST_FILTER,
project_id=PROJECT_ID,
)
notification_channel_enabled.enabled = False # pylint: disable=no-member
mask = FieldMask(paths=['enabled'])
mock_channel_client.return_value.update_notification_channel.assert_called_once_with(
request=dict(notification_channel=notification_channel_enabled,
update_mask=mask),
retry=DEFAULT,
timeout=DEFAULT,
metadata=(),
)
def test_execute(self, mock_hook):
operator = StackdriverListNotificationChannelsOperator(task_id=TEST_TASK_ID, filter_=TEST_FILTER)
mock_hook.return_value.list_notification_channels.return_value = [
NotificationChannel(name="test-123")
]
result = operator.execute(None)
mock_hook.return_value.list_notification_channels.assert_called_once_with(
project_id=None,
filter_=TEST_FILTER,
format_=None,
order_by=None,
page_size=None,
retry=DEFAULT,
timeout=DEFAULT,
metadata=None,
)
assert [
{
'description': '',
'display_name': '',
'labels': {},
'name': 'test-123',
'type_': '',
'user_labels': {},
'verification_status': 0,
}
] == result
def test_stackdriver_upsert_channel(self, mock_channel_client,
mock_get_creds_and_project_id):
hook = stackdriver.StackdriverHook()
existing_notification_channel = NotificationChannel(
**TEST_NOTIFICATION_CHANNEL_1)
notification_channel_to_be_created = NotificationChannel(
**TEST_NOTIFICATION_CHANNEL_2)
mock_channel_client.return_value.list_notification_channels.return_value = [
existing_notification_channel
]
hook.upsert_channel(
channels=json.dumps({
"channels":
[TEST_NOTIFICATION_CHANNEL_1, TEST_NOTIFICATION_CHANNEL_2]
}),
project_id=PROJECT_ID,
)
mock_channel_client.return_value.list_notification_channels.assert_called_once_with(
request=dict(name=f'projects/{PROJECT_ID}',
filter=None,
order_by=None,
page_size=None),
retry=DEFAULT,
timeout=DEFAULT,
metadata=(),
)
mock_channel_client.return_value.update_notification_channel.assert_called_once_with(
request=dict(notification_channel=existing_notification_channel),
retry=DEFAULT,
timeout=DEFAULT,
metadata=(),
)
notification_channel_to_be_created.name = None
mock_channel_client.return_value.create_notification_channel.assert_called_once_with(
request=dict(
name=f'projects/{PROJECT_ID}',
notification_channel=notification_channel_to_be_created),
retry=DEFAULT,
timeout=DEFAULT,
metadata=(),
)
def test_execute(self, mock_hook):
operator = StackdriverListNotificationChannelsOperator(
task_id=TEST_TASK_ID, filter_=TEST_FILTER)
mock_hook.return_value.list_notification_channels.return_value = [
NotificationChannel(name="test-123")
]
result = operator.execute(None)
mock_hook.return_value.list_notification_channels.assert_called_once_with(
project_id=None,
filter_=TEST_FILTER,
format_=None,
order_by=None,
page_size=None,
retry=DEFAULT,
timeout=DEFAULT,
metadata=None,
)
# Depending on the version of google-apitools installed we might receive the response either with or
# without mutation_records.
assert result in [
[{
'description': '',
'display_name': '',
'labels': {},
'name': 'test-123',
'type_': '',
'user_labels': {},
'verification_status': 0,
}],
[{
'description': '',
'display_name': '',
'labels': {},
'mutation_records': [],
'name': 'test-123',
'type_': '',
'user_labels': {},
'verification_status': 0,
}],
]
def upsert_channel(
self,
channels: str,
project_id: str,
retry: Optional[str] = DEFAULT,
timeout: Optional[float] = None,
metadata: Sequence[Tuple[str, str]] = (),
) -> dict:
"""
Creates a new notification or updates an existing notification channel
identified the name field in the alerts parameter.
:param channels: A JSON string or file that specifies all the alerts that needs
to be either created or updated. For more details, see
https://cloud.google.com/monitoring/api/ref_v3/rest/v3/projects.notificationChannels.
(templated)
:param project_id: The project in which notification channels needs to be created/updated.
:param retry: A retry object used to retry requests. If ``None`` is
specified, requests will be retried using a default configuration.
:param timeout: The amount of time, in seconds, to wait
for the request to complete. Note that if ``retry`` is
specified, the timeout applies to each individual attempt.
:param metadata: Additional metadata that is provided to the method.
"""
channel_client = self._get_channel_client()
record = json.loads(channels)
existing_channels = [
channel["name"] for channel in self.list_notification_channels(
project_id=project_id, format_="dict")
]
channels_list = []
channel_name_map = {}
for channel in record["channels"]:
channels_list.append(NotificationChannel(**channel))
for channel in channels_list:
channel.verification_status = (
monitoring_v3.NotificationChannel.VerificationStatus.
VERIFICATION_STATUS_UNSPECIFIED)
if channel.name in existing_channels:
channel_client.update_notification_channel(
request={'notification_channel': channel},
retry=retry,
timeout=timeout,
metadata=metadata,
)
else:
old_name = channel.name
channel.name = None
new_channel = channel_client.create_notification_channel(
request={
'name': f'projects/{project_id}',
'notification_channel': channel
},
retry=retry,
timeout=timeout,
metadata=metadata,
)
channel_name_map[old_name] = new_channel.name
return channel_name_map
def upsert_alert(
self,
alerts: str,
project_id: str = PROVIDE_PROJECT_ID,
retry: Optional[str] = DEFAULT,
timeout: Optional[float] = None,
metadata: Sequence[Tuple[str, str]] = (),
) -> None:
"""
Creates a new alert or updates an existing policy identified
the name field in the alerts parameter.
:param project_id: The project in which alert needs to be created/updated.
:param alerts: A JSON string or file that specifies all the alerts that needs
to be either created or updated. For more details, see
https://cloud.google.com/monitoring/api/ref_v3/rest/v3/projects.alertPolicies#AlertPolicy.
(templated)
:param retry: A retry object used to retry requests. If ``None`` is
specified, requests will be retried using a default configuration.
:param timeout: The amount of time, in seconds, to wait
for the request to complete. Note that if ``retry`` is
specified, the timeout applies to each individual attempt.
:param metadata: Additional metadata that is provided to the method.
"""
policy_client = self._get_policy_client()
channel_client = self._get_channel_client()
record = json.loads(alerts)
existing_policies = [
policy['name']
for policy in self.list_alert_policies(project_id=project_id,
format_='dict')
]
existing_channels = [
channel['name'] for channel in self.list_notification_channels(
project_id=project_id, format_='dict')
]
policies_ = []
channels = []
for channel in record.get("channels", []):
channels.append(NotificationChannel(**channel))
for policy in record.get("policies", []):
policies_.append(AlertPolicy(**policy))
channel_name_map = {}
for channel in channels:
channel.verification_status = (
monitoring_v3.NotificationChannel.VerificationStatus.
VERIFICATION_STATUS_UNSPECIFIED)
if channel.name in existing_channels:
channel_client.update_notification_channel(
request={'notification_channel': channel},
retry=retry,
timeout=timeout,
metadata=metadata,
)
else:
old_name = channel.name
channel.name = None
new_channel = channel_client.create_notification_channel(
request={
'name': f'projects/{project_id}',
'notification_channel': channel
},
retry=retry,
timeout=timeout,
metadata=metadata,
)
channel_name_map[old_name] = new_channel.name
for policy in policies_:
policy.creation_record = None
policy.mutation_record = None
for i, channel in enumerate(policy.notification_channels):
new_channel = channel_name_map.get(channel)
if new_channel:
policy.notification_channels[i] = new_channel
if policy.name in existing_policies:
try:
policy_client.update_alert_policy(
request={'alert_policy': policy},
retry=retry,
timeout=timeout,
metadata=metadata,
)
except InvalidArgument:
pass
else:
policy.name = None
for condition in policy.conditions:
condition.name = None
policy_client.create_alert_policy(
request={
'name': f'projects/{project_id}',
'alert_policy': policy
},
retry=retry,
timeout=timeout,
metadata=metadata,
)