removed service bus and modified other changes.

priyansh-sumo · priyansh-sumo · commit c64b52f83003 · 2024-01-18T20:22:01.000+05:30
diff --git a/EventHubs/requirements.txt b/EventHubs/requirements.txt
@@ -4,6 +4,6 @@ azure-identity==1.15.0
 azure-core==1.29.6
 azure-mgmt-resource==23.0.1
 azure-mgmt-eventhub==11.0.0
-azure-servicebus==0.21.1
+azure-eventhub==5.11.5
 azure-mgmt-storage==21.1.0
 zure-mgmt-loganalytics==12.0.0
diff --git a/EventHubs/tests/baseeventhubtest.py b/EventHubs/tests/baseeventhubtest.py
@@ -1,17 +1,18 @@
 import os
 import sys
 import json
-import time
-from time import sleep
-from requests import Session
-from datetime import datetime
-from azure.servicebus import ServiceBusService
+from datetime import timedelta
 from azure.mgmt.resource import ResourceManagementClient
 from azure.mgmt.eventhub import EventHubManagementClient
+from azure.eventhub import EventHubProducerClient
+from azure.eventhub.exceptions import EventHubError
+from azure.monitor.query import LogsQueryClient, LogsQueryStatus
+from azure.mgmt.loganalytics import LogAnalyticsManagementClient
 
 sys.path.insert(0, '../../test_utils')
 from basetest import BaseTest
 
+
 class BaseEventHubTest(BaseTest):
 
     def setUp(self):
@@ -41,69 +42,6 @@ def get_resource(self, restype):
                 return item
         raise Exception("%s Resource Not Found" % (restype))
 
-    def get_row_count(self, query):
-        rows = self.table_service.query_entities(
-            self.log_table_name, filter="PartitionKey eq 'R2'", select='PartitionKey')
-
-        return len(rows.items)
-    
-    def wait_for_table_results(self, query):
-        max_retries = 50
-        while(max_retries > 0 and (not (self.table_service.exists(
-              self.log_table_name) and self.get_row_count(query) > 0))):
-            print("waiting for logs creation...", max_retries)
-            sleep(15)
-            max_retries -= 1
-
-    def insert_mock_logs_in_EventHub(self, filename):
-        print("Inserting fake logs in EventHub")
-        namespace_name = self.get_resource_name(self.event_hub_namespace_prefix, "Microsoft.EventHub/namespaces")
-
-        defaultauthorule_name = "RootManageSharedAccessKey"
-
-        eventhub_client = EventHubManagementClient(self.credentials,
-                                                   self.subscription_id)
-
-        ehkeys = eventhub_client.namespaces.list_keys(
-            self.RESOURCE_GROUP_NAME, namespace_name, defaultauthorule_name)
-
-        sbs = ServiceBusService(
-            namespace_name,
-            shared_access_key_name=defaultauthorule_name,
-            shared_access_key_value=ehkeys.primary_key,
-            request_session=Session()
-        )
-        mock_logs = json.load(open(filename))
-        print("inserting %s" % (mock_logs))
-        sbs.send_event(self.eventhub_name, json.dumps(mock_logs))
-
-        print("Event inserted")
-        
-    def insert_mock_metrics_in_EventHub(self, filename):
-        print("Inserting fake metrics in EventHub")
-        
-        defaultauthorule_name = "RootManageSharedAccessKey"
-        namespace_name = self.get_resource_name(self.event_hub_namespace_prefix, "Microsoft.EventHub/namespaces")
-        eventhub_client = EventHubManagementClient(self.azure_credential, self.subscription_id)
-        eventhub_keys = eventhub_client.namespaces.list_keys(self.RESOURCE_GROUP_NAME, namespace_name, defaultauthorule_name)
-        
-        sbs = ServiceBusService(
-            namespace_name,
-            shared_access_key_name=defaultauthorule_name,
-            shared_access_key_value=eventhub_keys.primary_key,
-            request_session=Session()
-        )
-        
-        with open(filename, 'r') as template_file_fd:
-            mock_logs = json.load(template_file_fd)
-            mock_logs = json.dumps(mock_logs)
-            mock_logs = mock_logs.replace("2018-03-07T14:23:51.991Z", datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ"))
-            mock_logs = mock_logs.replace("C088DC46", "%d-%s" % (1, str(int(time.time()))))
-
-        # print("inserting %s" % (mock_logs))
-        sbs.send_event(self.eventhub_name, mock_logs)
-        print("Event inserted")
-
     def _parse_template(self):
         template_path = os.path.join(os.path.abspath('..'), 'src',
                                      self.template_name)
@@ -118,3 +56,83 @@ def _parse_template(self):
         template_data["parameters"]["location"]["defaultValue"] = self.resourcegroup_location
 
         return template_data
+    
+    def send_event_data_list(self, event_hub_namespace_prefix, event_hub_name, event_data_list):
+        
+        defaultauthorule_name = "RootManageSharedAccessKey"
+        namespace_name = self.get_resource_name(event_hub_namespace_prefix, "Microsoft.EventHub/namespaces")
+        eventhub_client = EventHubManagementClient(self.azure_credential, self.subscription_id)
+        eventhub_keys = eventhub_client.namespaces.list_keys(self.RESOURCE_GROUP_NAME, namespace_name, defaultauthorule_name)
+       
+        producer = EventHubProducerClient.from_connection_string(
+            conn_str=eventhub_keys.primary_connection_string,
+            eventhub_name=event_hub_name
+        )
+        
+        with producer:
+            try:
+                producer.send_batch(event_data_list)
+            except ValueError:  # Size exceeds limit. This shouldn't happen if you make sure before hand.
+                print("Size of the event data list exceeds the size limit of a single send")
+            except EventHubError as eh_err:
+                print("Sending error: ", eh_err)
+
+        print("Event inserted")
+    
+    def fetchlogs(self, app_insights):
+        result = []
+        try:
+            client = LogsQueryClient(self.azure_credential)
+            query = f"app('{app_insights}').traces | where operation_Name == '{self.function_name}' | project operation_Id, timestamp, message, severityLevel"
+            response = client.query_workspace(self.get_Workspace_Id(), query, timespan=timedelta(hours=1))
+                
+            if response.status == LogsQueryStatus.FAILURE:
+                raise Exception(f"LogsQueryError: {response.message}")
+            elif response.status == LogsQueryStatus.PARTIAL:
+                data = response.partial_data
+                error = response.partial_error
+                print("partial_error: ", error)
+            elif response.status == LogsQueryStatus.SUCCESS:
+                data = response.tables
+
+            for table in data:
+                for row in table.rows:
+                    row_dict = {str(col): str(item) for col, item in zip(table.columns, row)}
+                    result.append(row_dict)
+        except Exception as e:
+            print("An unexpected error occurred during the test:")
+            print("Exception", e)
+
+        return result
+    
+    def get_resources(self, resource_group_name):
+        return self.resource_client.resources.list_by_resource_group(resource_group_name)
+    
+    def get_Workspace_Id(self):
+        workspace = self.get_resource('microsoft.operationalinsights/workspaces')
+        client = LogAnalyticsManagementClient(
+            credential=self.azure_credential,
+            subscription_id=self.subscription_id,
+        )
+
+        response = client.workspaces.get(
+            resource_group_name=self.RESOURCE_GROUP_NAME,
+            workspace_name=workspace.name,
+        )
+        return response.customer_id
+    
+    def filter_logs(self, logs, key, value):
+        return value in [d.get(key) for d in logs]
+    
+    def check_resource_count(self):
+        resource_count = len(list(self.get_resources(self.RESOURCE_GROUP_NAME)))
+        self.assertTrue(resource_count == self.expected_resource_count, f"resource count of resource group {self.RESOURCE_GROUP_NAME} differs from expected count : {resource_count}")
+
+    def check_success_log(self, logs):
+        self.assertTrue(self.filter_logs(logs, 'message', self.successful_sent_message))
+    
+    def check_error_log(self, logs):
+        self.assertTrue(not self.filter_logs(logs, 'severityLevel', '3'))
+
+    def check_warning_log(self, logs):
+        self.assertTrue(not self.filter_logs(logs, 'severityLevel', '2'))
diff --git a/EventHubs/tests/test_eventhub_logs.py b/EventHubs/tests/test_eventhub_logs.py
@@ -1,20 +1,20 @@
 import unittest
-import datetime
+from datetime import datetime
+import time
+import json
 from baseeventhubtest import BaseEventHubTest
+from azure.eventhub import EventData
 
 
 class TestEventHubLogs(BaseEventHubTest):
 
     def setUp(self):
         super(TestEventHubLogs, self).setUp()
         self.RESOURCE_GROUP_NAME = "TestEventHubLogs-%s" % (
-            datetime.datetime.now().strftime("%d-%m-%y-%H-%M-%S"))
-        self.STORAGE_ACCOUNT_NAME = "sumoapplogs"
+            datetime.now().strftime("%d-%m-%y-%H-%M-%S"))
         self.function_name_prefix = "EventHubs_Logs"
         self.template_name = 'azuredeploy_logs.json'
         self.event_hub_namespace_prefix = "SumoAzureLogsNamespace"
-        self.log_table_name = "AzureWebJobsHostLogs%d%02d" % (
-            datetime.datetime.now().year, datetime.datetime.now().month)
         self.eventhub_name = 'insights-operational-logs'
 
     def test_pipeline(self):
@@ -24,24 +24,21 @@ def test_pipeline(self):
         self.assertTrue(self.resource_group_exists(self.RESOURCE_GROUP_NAME))
         self.table_service = self.get_table_service()
         self.insert_mock_logs_in_EventHub('activity_log_fixtures.json')
-        self.check_error_logs()
-
-    def check_error_logs(self):
-        print("sleeping 1min for function execution")
-        query = "PartitionKey eq 'R2'"
-        self.wait_for_table_results(query)
-        self.assertTrue(self.get_row_count(query) > 0)
-
-        rows = self.table_service.query_entities(
-            self.log_table_name, filter=query)
-
-        haserr = False
-        for row in rows.items:
-            print("LogRow: ", row["FunctionName"], row["HasError"])
-            if row["FunctionName"].startswith(self.function_name_prefix) and row["HasError"]:
-                haserr = True
-
-        self.assertTrue(not haserr)
+        
+    def insert_mock_logs_in_EventHub(self, filename):
+        print("Inserting fake logs in EventHub")
+        
+        with open(filename, 'r') as template_file_fd:
+            mock_logs = json.load(template_file_fd)
+            mock_logs = json.dumps(mock_logs)
+            mock_logs = mock_logs.replace("2018-03-07T14:23:51.991Z", datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ"))
+            mock_logs = mock_logs.replace("C088DC46", "%d-%s" % (1, str(int(time.time()))))
+
+        event_data_list = [EventData(mock_logs)]
+        # print("inserting %s" % (mock_logs))
+        self.send_event_data_list(self.event_hub_namespace_prefix, self.eventhub_name, event_data_list)
+
+        print("Event inserted")
 
 
 if __name__ == '__main__':
diff --git a/EventHubs/tests/test_eventhub_metrics.py b/EventHubs/tests/test_eventhub_metrics.py
@@ -1,11 +1,10 @@
-import sys
 import unittest
 import time
-from io import StringIO
-from datetime import timedelta, datetime
+import json
+from datetime import datetime
 from baseeventhubtest import BaseEventHubTest
-from azure.monitor.query import LogsQueryClient, LogsQueryStatus
-from azure.mgmt.loganalytics import LogAnalyticsManagementClient
+from azure.eventhub import EventData
+
 
 class TestEventHubMetrics(BaseEventHubTest):
 
@@ -16,75 +15,35 @@ def setUp(self):
         self.event_hub_namespace_prefix = "SMNamespace"
         self.eventhub_name = "insights-metrics-pt1m"
         self.function_name = "EventHubs_Metrics"
-        self.resource_count = 8
+        self.successful_sent_message = 'Sent all metric data to Sumo. Exit now.'
+        self.expected_resource_count = 7
         
     def test_pipeline(self):
         self.create_resource_group()
         self.deploy_template()
         self.assertTrue(self.resource_group_exists(self.RESOURCE_GROUP_NAME)) 
         self.insert_mock_metrics_in_EventHub('metrics_fixtures.json')
         time.sleep(300)  # Due to latency, logs are available after few mins.
-        self.get_resource_count()
+        self.check_resource_count()
         app_insights = self.get_resource('Microsoft.Insights/components')
         captured_output = self.fetchlogs(app_insights.name)
         self.check_success_log(captured_output)
         self.check_error_log(captured_output)
         self.check_warning_log(captured_output)
     
-    def fetchlogs(self, app_insights):
-        result = []
-        try:
-            client = LogsQueryClient(self.azure_credential)
-            query = f"app('{app_insights}').traces | where operation_Name == '{self.function_name}' | project operation_Id, timestamp, message, severityLevel"
-            response = client.query_workspace(self.get_Workspace_Id(), query, timespan=timedelta(hours=1))
-                
-            if response.status == LogsQueryStatus.PARTIAL:
-                data = response.partial_data
-            elif response.status == LogsQueryStatus.SUCCESS:
-                data = response.tables
-
-            for table in data:
-                for row in table.rows:
-                    row_dict = {str(col): str(item) for col, item in zip(table.columns, row)}
-                    result.append(row_dict)
-        except Exception as e:
-            print("An unexpected error occurred during the test:")
-            print("Exception", e)
-
-        return result
-    
-    def get_resource_count(self):
-        return len(list(self.resource_client.resources.list_by_resource_group(self.RESOURCE_GROUP_NAME)))
-    
-    def filter_logs(self, logs, key, value):
-        return value in [d.get(key) for d in logs]
-    
-    def check_resource_count(self):
-        self.assertTrue(self.get_resource_count() == self.resource_count)
-
-    def check_success_log(self, logs):
-        successful_sent_message = 'Sent all metric data to Sumo. Exit now.'
-        self.assertTrue(self.filter_logs(logs, 'message', successful_sent_message))
-
-    def check_error_log(self, logs):
-        self.assertTrue(not self.filter_logs(logs, 'severityLevel', '3'))
-
-    def check_warning_log(self, logs):
-        self.assertTrue(not self.filter_logs(logs, 'severityLevel', '2'))
+    def insert_mock_metrics_in_EventHub(self, filename):
+        print("Inserting fake metrics in EventHub")
+        
+        with open(filename, 'r') as template_file_fd:
+            mock_logs = json.load(template_file_fd)
+            mock_logs = json.dumps(mock_logs)
+            mock_logs = mock_logs.replace("2018-03-07T14:23:51.991Z", datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ"))
+            mock_logs = mock_logs.replace("C088DC46", "%d-%s" % (1, str(int(time.time()))))
 
-    def get_Workspace_Id(self):
-        workspace = self.get_resource('microsoft.operationalinsights/workspaces')
-        client = LogAnalyticsManagementClient(
-            credential=self.azure_credential,
-            subscription_id=self.subscription_id,
-        )
+        event_data_list = [EventData(mock_logs)]
+        # print("inserting %s" % (mock_logs))
+        self.send_event_data_list(self.event_hub_namespace_prefix, self.eventhub_name, event_data_list)
 
-        response = client.workspaces.get(
-            resource_group_name=self.RESOURCE_GROUP_NAME,
-            workspace_name=workspace.name,
-        )
-        return response.customer_id
-    
 
 if __name__ == '__main__':
     unittest.main()
