Update tests

Author: samwelkanda

tests/unit/test_async_graph_transport.py

@@ -0,0 +1,18 @@
+import pytest
+from kiota_http.kiota_client_factory import KiotaClientFactory
+
+from msgraph_core._enums import FeatureUsageFlag
+from msgraph_core.middleware import AsyncGraphTransport, GraphRequestContext
+
+
+def test_set_request_context_and_feature_usage(mock_request, mock_transport):
+    middleware = KiotaClientFactory.get_default_middleware()
+    pipeline = KiotaClientFactory.create_middleware_pipeline(middleware, mock_transport)
+    transport = AsyncGraphTransport(mock_transport, pipeline)
+    transport.set_request_context_and_feature_usage(mock_request)
+
+    assert hasattr(mock_request, 'context')
+    assert isinstance(mock_request.context, GraphRequestContext)
+    assert mock_request.context.feature_usage == hex(
+        FeatureUsageFlag.RETRY_HANDLER_ENABLED | FeatureUsageFlag.REDIRECT_HANDLER_ENABLED
+    )

tests/unit/test_graph_client_factory.py

@@ -4,18 +4,18 @@
 # ------------------------------------
 import httpx
 import pytest
-from kiota_http.middleware import AsyncKiotaTransport, MiddlewarePipeline, RedirectHandler
+from kiota_http.middleware import MiddlewarePipeline, RedirectHandler
 
 from msgraph_core import APIVersion, GraphClientFactory, NationalClouds
-from msgraph_core.middleware.telemetry import GraphTelemetryHandler
+from msgraph_core.middleware import AsyncGraphTransport, GraphTelemetryHandler
 
 
 def test_create_with_default_middleware():
     """Test creation of GraphClient using default middleware"""
     client = GraphClientFactory.create_with_default_middleware()
 
     assert isinstance(client, httpx.AsyncClient)
-    assert isinstance(client._transport, AsyncKiotaTransport)
+    assert isinstance(client._transport, AsyncGraphTransport)
     pipeline = client._transport.pipeline
     assert isinstance(pipeline, MiddlewarePipeline)
     assert isinstance(pipeline._first_middleware, RedirectHandler)
@@ -30,7 +30,7 @@ def test_create_with_custom_middleware():
     client = GraphClientFactory.create_with_custom_middleware(middleware=middleware)
 
     assert isinstance(client, httpx.AsyncClient)
-    assert isinstance(client._transport, AsyncKiotaTransport)
+    assert isinstance(client._transport, AsyncGraphTransport)
     pipeline = client._transport.pipeline
     assert isinstance(pipeline._first_middleware, GraphTelemetryHandler)
 

tests/unit/test_graph_telemetry_handler.py

@@ -10,22 +10,11 @@
 import pytest
 
 from msgraph_core import SDK_VERSION, APIVersion, NationalClouds
-from msgraph_core._enums import FeatureUsageFlag
 from msgraph_core.middleware import GraphRequestContext, GraphTelemetryHandler
 
 BASE_URL = NationalClouds.Global + '/' + APIVersion.v1
 
 
-def test_set_request_context_and_feature_usage(mock_request, mock_transport):
-    telemetry_handler = GraphTelemetryHandler()
-    telemetry_handler.set_request_context_and_feature_usage(mock_request, mock_transport)
-
-    assert hasattr(mock_request, 'context')
-    assert mock_request.context.feature_usage == hex(
-        FeatureUsageFlag.RETRY_HANDLER_ENABLED | FeatureUsageFlag.REDIRECT_HANDLER_ENABLED
-    )
-
-
 def test_is_graph_url(mock_graph_request):
     """
     Test method that checks whether a request url is a graph endpoint