test_webhook.py
1 import pytest 2 3 from mlflow.entities.webhook import ( 4 Webhook, 5 WebhookAction, 6 WebhookEntity, 7 WebhookEvent, 8 WebhookStatus, 9 WebhookTestResult, 10 ) 11 from mlflow.exceptions import MlflowException 12 from mlflow.protos.webhooks_pb2 import WebhookAction as ProtoWebhookAction 13 from mlflow.protos.webhooks_pb2 import WebhookEntity as ProtoWebhookEntity 14 from mlflow.protos.webhooks_pb2 import WebhookStatus as ProtoWebhookStatus 15 from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME 16 17 18 @pytest.mark.parametrize( 19 ("proto_status", "status_enum"), 20 [ 21 (ProtoWebhookStatus.ACTIVE, WebhookStatus.ACTIVE), 22 (ProtoWebhookStatus.DISABLED, WebhookStatus.DISABLED), 23 ], 24 ) 25 def test_webhook_status_proto_conversion(proto_status, status_enum): 26 assert WebhookStatus.from_proto(proto_status) == status_enum 27 assert status_enum.to_proto() == proto_status 28 29 30 @pytest.mark.parametrize( 31 ("entity_enum", "proto_entity"), 32 [ 33 (WebhookEntity.REGISTERED_MODEL, ProtoWebhookEntity.REGISTERED_MODEL), 34 (WebhookEntity.MODEL_VERSION, ProtoWebhookEntity.MODEL_VERSION), 35 (WebhookEntity.MODEL_VERSION_TAG, ProtoWebhookEntity.MODEL_VERSION_TAG), 36 (WebhookEntity.MODEL_VERSION_ALIAS, ProtoWebhookEntity.MODEL_VERSION_ALIAS), 37 ], 38 ) 39 def test_webhook_entity_proto_conversion(entity_enum, proto_entity): 40 assert WebhookEntity.from_proto(proto_entity) == entity_enum 41 assert entity_enum.to_proto() == proto_entity 42 43 44 @pytest.mark.parametrize( 45 ("action_enum", "proto_action"), 46 [ 47 (WebhookAction.CREATED, ProtoWebhookAction.CREATED), 48 (WebhookAction.UPDATED, ProtoWebhookAction.UPDATED), 49 (WebhookAction.DELETED, ProtoWebhookAction.DELETED), 50 (WebhookAction.SET, ProtoWebhookAction.SET), 51 ], 52 ) 53 def test_webhook_action_proto_conversion(action_enum, proto_action): 54 assert WebhookAction.from_proto(proto_action) == action_enum 55 assert action_enum.to_proto() == proto_action 56 57 58 def test_webhook_event_creation(): 59 event = WebhookEvent(WebhookEntity.REGISTERED_MODEL, WebhookAction.CREATED) 60 assert event.entity == WebhookEntity.REGISTERED_MODEL 61 assert event.action == WebhookAction.CREATED 62 63 64 def test_webhook_event_from_string(): 65 event = WebhookEvent("registered_model", "created") 66 assert event.entity == WebhookEntity.REGISTERED_MODEL 67 assert event.action == WebhookAction.CREATED 68 69 70 def test_webhook_event_invalid_combination(): 71 with pytest.raises( 72 MlflowException, match="Invalid action 'updated' for entity 'model_version_tag'" 73 ): 74 WebhookEvent(WebhookEntity.MODEL_VERSION_TAG, WebhookAction.UPDATED) 75 76 77 def test_webhook_event_from_str(): 78 event = WebhookEvent.from_str("registered_model.created") 79 assert event.entity == WebhookEntity.REGISTERED_MODEL 80 assert event.action == WebhookAction.CREATED 81 82 83 def test_webhook_event_from_str_invalid_format(): 84 with pytest.raises(MlflowException, match="Invalid event string format"): 85 WebhookEvent.from_str("invalid_format") 86 87 88 def test_webhook_event_to_str(): 89 event = WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED) 90 assert str(event) == "model_version.created" 91 92 93 def test_webhook_event_proto_conversion(): 94 event = WebhookEvent(WebhookEntity.REGISTERED_MODEL, WebhookAction.CREATED) 95 proto_event = event.to_proto() 96 event_from_proto = WebhookEvent.from_proto(proto_event) 97 assert event_from_proto.entity == event.entity 98 assert event_from_proto.action == event.action 99 100 101 def test_webhook_event_equality(): 102 event1 = WebhookEvent(WebhookEntity.REGISTERED_MODEL, WebhookAction.CREATED) 103 event2 = WebhookEvent(WebhookEntity.REGISTERED_MODEL, WebhookAction.CREATED) 104 event3 = WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED) 105 106 assert event1 == event2 107 assert event1 != event3 108 assert hash(event1) == hash(event2) 109 assert hash(event1) != hash(event3) 110 111 112 def test_webhook_event_invalid_entity_action_combination(): 113 with pytest.raises( 114 MlflowException, match="Invalid action 'deleted' for entity 'registered_model'" 115 ): 116 WebhookEvent(WebhookEntity.REGISTERED_MODEL, WebhookAction.DELETED) 117 118 119 def test_webhook_proto_conversion(): 120 events = [ 121 WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED), 122 WebhookEvent(WebhookEntity.MODEL_VERSION_ALIAS, WebhookAction.CREATED), 123 ] 124 webhook = Webhook( 125 webhook_id="webhook123", 126 name="Test Webhook", 127 url="https://example.com/webhook", 128 events=events, 129 description="Test webhook description", 130 status=WebhookStatus.ACTIVE, 131 secret="my-secret", 132 creation_timestamp=1234567890, 133 last_updated_timestamp=1234567900, 134 ) 135 proto_webhook = webhook.to_proto() 136 webhook_from_proto = Webhook.from_proto(proto_webhook) 137 assert webhook_from_proto.webhook_id == webhook.webhook_id 138 assert webhook_from_proto.name == webhook.name 139 assert webhook_from_proto.url == webhook.url 140 assert webhook_from_proto.events == webhook.events 141 assert webhook_from_proto.description == webhook.description 142 assert webhook_from_proto.status == webhook.status 143 assert webhook_from_proto.creation_timestamp == webhook.creation_timestamp 144 assert webhook_from_proto.last_updated_timestamp == webhook.last_updated_timestamp 145 assert webhook_from_proto.workspace == DEFAULT_WORKSPACE_NAME 146 147 148 def test_webhook_workspace_defaults_to_default_workspace(): 149 events = [WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)] 150 webhook = Webhook( 151 webhook_id="webhook123", 152 name="Test Webhook", 153 url="https://example.com/webhook", 154 events=events, 155 creation_timestamp=1234567890, 156 last_updated_timestamp=1234567900, 157 workspace=None, 158 ) 159 assert webhook.workspace == DEFAULT_WORKSPACE_NAME 160 161 162 def test_webhook_workspace_is_preserved(): 163 events = [WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)] 164 webhook = Webhook( 165 webhook_id="webhook123", 166 name="Team Webhook", 167 url="https://example.com/webhook", 168 events=events, 169 creation_timestamp=1234567890, 170 last_updated_timestamp=1234567900, 171 workspace="team-a", 172 ) 173 assert webhook.workspace == "team-a" 174 175 176 def test_webhook_no_secret_in_repr(): 177 events = [WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)] 178 webhook = Webhook( 179 webhook_id="webhook123", 180 name="Test Webhook", 181 url="https://example.com/webhook", 182 events=events, 183 creation_timestamp=1234567890, 184 last_updated_timestamp=1234567900, 185 description="Test webhook description", 186 status=WebhookStatus.ACTIVE, 187 secret="my-secret", 188 ) 189 assert "my-secret" not in repr(webhook) 190 191 192 def test_webhook_invalid_events(): 193 with pytest.raises(MlflowException, match="Webhook events cannot be empty"): 194 Webhook( 195 webhook_id="webhook123", 196 name="Test Webhook", 197 url="https://example.com/webhook", 198 events=[], 199 creation_timestamp=1234567890, 200 last_updated_timestamp=1234567900, 201 ) 202 203 204 def test_webhook_test_result(): 205 # Test successful result 206 result = WebhookTestResult( 207 success=True, 208 response_status=200, 209 response_body='{"status": "ok"}', 210 ) 211 assert result.success is True 212 assert result.response_status == 200 213 assert result.response_body == '{"status": "ok"}' 214 assert result.error_message is None 215 216 # Test failed result 217 result = WebhookTestResult( 218 success=False, 219 response_status=500, 220 error_message="Internal server error", 221 ) 222 assert result.success is False 223 assert result.response_status == 500 224 assert result.error_message == "Internal server error" 225 assert result.response_body is None