Add more tests, including asyncio_ensure_future wrapper

Author: lmazuel

sdk/core/azure-core/azure/core/polling/_async_poller.py

@@ -128,6 +128,11 @@ def __init__(
 
         self._polling_method.initialize(client, initial_response, deserialization_callback)
 
+    def polling_method(self) -> AsyncPollingMethod[PollingReturnType]:
+        """Return the polling method associated to this poller.
+        """
+        return self._polling_method
+
     def continuation_token(self) -> str:
         """Return a continuation token that allows to restart the poller later.
 

sdk/core/azure-core/azure/core/polling/_poller.py

@@ -196,6 +196,12 @@ def _start(self):
                 call(self._polling_method)
             callbacks, self._callbacks = self._callbacks, []
 
+    def polling_method(self):
+        # type: () -> PollingMethod[PollingReturnType]
+        """Return the polling method associated to this poller.
+        """
+        return self._polling_method
+
     def continuation_token(self):
         # type: () -> str
         """Return a continuation token that allows to restart the poller later.

sdk/core/azure-core/tests/azure_core_asynctests/test_polling.py

@@ -23,6 +23,7 @@
 # THE SOFTWARE.
 #
 #--------------------------------------------------------------------------
+import asyncio
 import time
 try:
     from unittest import mock
@@ -89,6 +90,7 @@ async def run(self):
         """Empty run, no polling.
         """
         self._finished = True
+        await asyncio.sleep(self._sleep) # Give me time to add callbacks!
 
     def status(self):
         """Return the current status as a string.
@@ -129,39 +131,53 @@ def deserialization_callback(response):
 
     method = AsyncNoPolling()
 
-    poller = AsyncLROPoller(client, initial_response, deserialization_callback, method)
+    raw_poller = AsyncLROPoller(client, initial_response, deserialization_callback, method)
+    poller = asyncio.ensure_future(raw_poller.result())
+
+    done_cb = mock.MagicMock()
+    poller.add_done_callback(done_cb)
 
-    result = await poller.result()
+    result = await poller
     assert poller.done()
     assert result == "Treated: "+initial_response
-    assert poller.status() == "succeeded"
+    assert raw_poller.status() == "succeeded"
+    assert raw_poller.polling_method() is method
+    done_cb.assert_called_once_with(poller)
 
     # Test with a basic Model
     poller = AsyncLROPoller(client, initial_response, Model, method)
     assert poller._polling_method._deserialization_callback == Model.deserialize
 
     # Test poller that method do a run
     method = PollingTwoSteps(sleep=1)
-    poller = AsyncLROPoller(client, initial_response, deserialization_callback, method)
+    raw_poller = AsyncLROPoller(client, initial_response, deserialization_callback, method)
+    poller = asyncio.ensure_future(raw_poller.result())
+
+    done_cb = mock.MagicMock()
+    done_cb2 = mock.MagicMock()
+    poller.add_done_callback(done_cb)
+    poller.remove_done_callback(done_cb2)
 
-    result = await poller.result()
+    result = await poller
     assert result == "Treated: "+initial_response
-    assert poller.status() == "succeeded"
+    assert raw_poller.status() == "succeeded"
+    done_cb.assert_called_once_with(poller)
+    done_cb2.assert_not_called()
 
     # Test continuation token
-    cont_token = poller.continuation_token()
+    cont_token = raw_poller.continuation_token()
 
     method = PollingTwoSteps(sleep=1)
     new_poller = AsyncLROPoller.from_continuation_token(
         continuation_token=cont_token,
         client=client,
         initial_response=initial_response,
-        deserialization_callback=Model,
+        deserialization_callback=deserialization_callback,
         polling_method=method
     )
-    result = await poller.result()
+    result = await new_poller.result()
     assert result == "Treated: "+initial_response
-    assert poller.status() == "succeeded"
+    assert new_poller.status() == "succeeded"
 
 
 @pytest.mark.asyncio

sdk/core/azure-core/tests/test_polling.py

@@ -163,6 +163,7 @@ def deserialization_callback(response):
     assert poller.done()
     assert result == "Treated: "+initial_response
     assert poller.status() == "succeeded"
+    assert poller.polling_method() is method
     done_cb.assert_called_once_with(method)
 
     # Test with a basic Model
@@ -196,12 +197,12 @@ def deserialization_callback(response):
         continuation_token=cont_token,
         client=client,
         initial_response=initial_response,
-        deserialization_callback=Model,
+        deserialization_callback=deserialization_callback,
         polling_method=method
     )
-    result = poller.result()
+    result = new_poller.result()
     assert result == "Treated: "+initial_response
-    assert poller.status() == "succeeded"
+    assert new_poller.status() == "succeeded"
 
 
 def test_broken_poller(client):