[text analytics] fix stats for analyze and healthcare

sdk/textanalytics/azure-ai-textanalytics/azure/ai/textanalytics/_policies.py

@@ -7,20 +7,26 @@
 from azure.core.pipeline.policies import ContentDecodePolicy
 from azure.core.pipeline.policies import SansIOHTTPPolicy
 from ._models import TextDocumentBatchStatistics
+from ._lro import _FINISHED
 
 
 class TextAnalyticsResponseHookPolicy(SansIOHTTPPolicy):
     def __init__(self, **kwargs):
         self._response_callback = kwargs.get("raw_response_hook")
+        self._is_lro = None
         super(TextAnalyticsResponseHookPolicy, self).__init__()
 
     def on_request(self, request):
         self._response_callback = request.context.options.pop("raw_response_hook", self._response_callback)
 
     def on_response(self, request, response):
-
+        if self._is_lro is None:
+            # determine LRO based off of initial response. If 202, we say it's an LRO
+            self._is_lro = response.http_response.status_code == 202
         if self._response_callback:
             data = ContentDecodePolicy.deserialize_from_http_generics(response.http_response)
+            if self._is_lro and (not data or data.get("status") not in _FINISHED):
+                return
             if data:
                 statistics = data.get("statistics", None)
                 model_version = data.get("modelVersion", None)

sdk/textanalytics/azure-ai-textanalytics/tests/test_analyze.py

@@ -354,8 +354,23 @@ def test_out_of_order_ids_multiple_tasks(self, client):
     def test_show_stats_and_model_version_multiple_tasks(self, client):
 
         def callback(resp):
-            if resp.raw_response:
-                a = "b"
+            assert resp.raw_response
+            tasks = resp.raw_response['tasks']
+            assert tasks['completed'] == 5
+            assert tasks['inProgress'] == 0
+            assert tasks['failed'] == 0
+            assert tasks['total'] == 5
+            num_tasks = 0
+            for key, task in tasks.items():
+                if "Tasks" in key:
+                    num_tasks += 1
+                    assert len(task) == 1
+                    task_stats = task[0]['results']['statistics']
+                    assert task_stats['documentsCount'] == 4
+                    assert task_stats['validDocumentsCount'] == 4
+                    assert task_stats['erroneousDocumentsCount'] == 0
+                    assert task_stats['transactionsCount'] == 4
+            assert num_tasks == 5
 
         docs = [{"id": "56", "text": ":)"},
                 {"id": "0", "text": ":("},

sdk/textanalytics/azure-ai-textanalytics/tests/test_analyze_async.py

@@ -391,6 +391,25 @@ async def test_show_stats_and_model_version_multiple_tasks(self, client):
                 {"id": "19", "text": ":P"},
                 {"id": "1", "text": ":D"}]
 
+        def callback(resp):
+            assert resp.raw_response
+            tasks = resp.raw_response['tasks']
+            assert tasks['completed'] == 5
+            assert tasks['inProgress'] == 0
+            assert tasks['failed'] == 0
+            assert tasks['total'] == 5
+            num_tasks = 0
+            for key, task in tasks.items():
+                if "Tasks" in key:
+                    num_tasks += 1
+                    assert len(task) == 1
+                    task_stats = task[0]['results']['statistics']
+                    assert task_stats['documentsCount'] == 4
+                    assert task_stats['validDocumentsCount'] == 4
+                    assert task_stats['erroneousDocumentsCount'] == 0
+                    assert task_stats['transactionsCount'] == 4
+            assert num_tasks == 5
+
         async with client:
             response = await (await client.begin_analyze_actions(
                 docs,
@@ -402,7 +421,8 @@ async def test_show_stats_and_model_version_multiple_tasks(self, client):
                     AnalyzeSentimentAction(model_version="latest")
                 ],
                 show_stats=True,
-                polling_interval=self._interval()
+                polling_interval=self._interval(),
+                raw_response_hook=callback,
             )).result()
 
             action_results = []

sdk/textanalytics/azure-ai-textanalytics/tests/test_analyze_healthcare.py

@@ -132,11 +132,20 @@ def test_show_stats_and_model_version(self, client):
                 {"id": "19", "text": ":P"},
                 {"id": "1", "text": ":D"}]
 
+        def callback(resp):
+            assert resp.raw_response
+            stats = resp.raw_response['results']['statistics']
+            assert stats['documentsCount'] == 5
+            assert stats['validDocumentsCount'] == 4
+            assert stats['erroneousDocumentsCount'] == 1
+            assert stats['transactionsCount'] == 4
+
         response = client.begin_analyze_healthcare_entities(
             docs,
             show_stats=True,
             model_version="2021-01-11",
-            polling_interval=self._interval()
+            polling_interval=self._interval(),
+            raw_response_hook = callback,
         ).result()
 
         num_error = 0

sdk/textanalytics/azure-ai-textanalytics/tests/test_analyze_healthcare_async.py

@@ -158,12 +158,21 @@ async def test_show_stats_and_model_version(self, client):
                 {"id": "19", "text": ":P"},
                 {"id": "1", "text": ":D"}]
 
+        def callback(resp):
+            assert resp.raw_response
+            stats = resp.raw_response['results']['statistics']
+            assert stats['documentsCount'] == 5
+            assert stats['validDocumentsCount'] == 4
+            assert stats['erroneousDocumentsCount'] == 1
+            assert stats['transactionsCount'] == 4
+
         async with client:
             response = await (await client.begin_analyze_healthcare_entities(
                 docs,
                 show_stats=True,
                 model_version="2021-01-11",
-                polling_interval=self._interval()
+                polling_interval=self._interval(),
+                raw_response_hook=callback,
             )).result()
 
         assert response

tools/azure-sdk-tools/devtools_testutils/cognitiveservices_testcase.py

@@ -62,7 +62,7 @@ def create_resource(self, name, **kwargs):
         if self.is_live:
             self.client = self.create_mgmt_client(CognitiveServicesManagementClient)
             group = self._get_resource_group(**kwargs)
-            cogsci_account = self.client.accounts.create(
+            cogsci_account = self.client.accounts.begin_create(
                 group.name,
                 name,
                 account={
@@ -71,7 +71,7 @@ def create_resource(self, name, **kwargs):
                     "kind": self.kind,
                     "properties": {"custom_sub_domain_name": self.custom_subdomain_name},
                 },
-            )
+            ).result()
             time.sleep(10)  # it takes a few seconds to create a cognitive services account
             self.resource = cogsci_account
             self.cogsci_key = self.client.accounts.list_keys(group.name, name).key1
@@ -114,7 +114,7 @@ def create_resource(self, name, **kwargs):
     def remove_resource(self, name, **kwargs):
         if self.is_live:
             group = self._get_resource_group(**kwargs)
-            self.client.accounts.delete(group.name, name)
+            self.client.accounts.begin_delete(group.name, name).wait()
 
     def _get_resource_group(self, **kwargs):
         try: