Change future value behavior of message state (Azure#22949)

* Change future value behavior

* more tests

* test active

Author: rakshith91

sdk/servicebus/azure-servicebus/azure/servicebus/_common/message.py

@@ -971,18 +971,19 @@ def dead_letter_source(self):
 
     @property
     def message_state(self):
-        # type: () -> ServiceBusMessageState
+        # type: () -> int
         """
         Defaults to Active. Represents the message state of the message. Can be Active, Deferred.
         or Scheduled.
 
-        :rtype: ~azure.servicebus.ServiceBusMessageState
+        :rtype: int
         """
         try:
             message_state = self._raw_amqp_message.annotations.get(MESSAGE_STATE_NAME)
-            if not message_state:
-                return ServiceBusMessageState.ACTIVE
-            return ServiceBusMessageState(message_state)
+            try:
+                return ServiceBusMessageState(message_state)
+            except ValueError:
+                return ServiceBusMessageState.ACTIVE if not message_state else message_state
         except AttributeError:
             return ServiceBusMessageState.ACTIVE
 

sdk/servicebus/azure-servicebus/tests/test_message.py

@@ -1,6 +1,6 @@
 import uamqp
 from datetime import datetime, timedelta
-from azure.servicebus import ServiceBusMessage, ServiceBusReceivedMessage
+from azure.servicebus import ServiceBusMessage, ServiceBusReceivedMessage, ServiceBusMessageState
 from azure.servicebus._common.constants import (
     _X_OPT_PARTITION_KEY,
     _X_OPT_VIA_PARTITION_KEY,
@@ -56,6 +56,52 @@ def test_servicebus_received_message_repr():
     assert "content_type=None, correlation_id=None, to=None, reply_to=None, reply_to_session_id=None, subject=None,"
     assert "partition_key=r_key, scheduled_enqueue_time_utc" in repr_str
 
+def test_servicebus_received_message_state():
+    uamqp_received_message = uamqp.message.Message(
+        body=b'data',
+        annotations={
+            b"x-opt-message-state": 3
+        },
+        properties=uamqp.message.MessageProperties()
+    )
+    received_message = ServiceBusReceivedMessage(uamqp_received_message, receiver=None)
+    assert received_message.message_state == 3
+
+    uamqp_received_message = uamqp.message.Message(
+        body=b'data',
+        annotations={
+            b"x-opt-message-state": 1
+        },
+        properties=uamqp.message.MessageProperties()
+    )
+    received_message = ServiceBusReceivedMessage(uamqp_received_message, receiver=None)
+    assert received_message.message_state == ServiceBusMessageState.DEFERRED
+
+    uamqp_received_message = uamqp.message.Message(
+        body=b'data',
+        annotations={
+        },
+        properties=uamqp.message.MessageProperties()
+    )
+    received_message = ServiceBusReceivedMessage(uamqp_received_message, receiver=None)
+    assert received_message.message_state == ServiceBusMessageState.ACTIVE
+
+    uamqp_received_message = uamqp.message.Message(
+        body=b'data',
+        properties=uamqp.message.MessageProperties()
+    )
+    received_message = ServiceBusReceivedMessage(uamqp_received_message, receiver=None)
+    assert received_message.message_state == ServiceBusMessageState.ACTIVE
+
+    uamqp_received_message = uamqp.message.Message(
+        body=b'data',
+        annotations={
+            b"x-opt-message-state": 0
+        },
+        properties=uamqp.message.MessageProperties()
+    )
+    received_message = ServiceBusReceivedMessage(uamqp_received_message, receiver=None)
+    assert received_message.message_state == ServiceBusMessageState.ACTIVE
 
 def test_servicebus_received_message_repr_with_props():
     uamqp_received_message = uamqp.message.Message(