WIP: try prometheus_client HTTP server instead of MicroService in test

Signed-off-by: Eero Tamminen <[email protected]>

Author: eero-t

tests/cores/mega/test_service_orchestrator_streaming.py

@@ -8,19 +8,11 @@
 
 import requests
 from fastapi.responses import StreamingResponse
+from prometheus_client import start_http_server
 
-from comps import (
-    MicroService,
-    ServiceOrchestrator,
-    ServiceRoleType,
-    ServiceType,
-    TextDoc,
-    opea_microservices,
-    register_microservice,
-)
-from comps.cores.proto.api_protocol import ChatCompletionRequest, ChatCompletionResponse
+from comps import ServiceOrchestrator, ServiceType, TextDoc, opea_microservices, register_microservice
 
-_MEGA_PORT = 8888
+_METRIC_PORT = 8000
 
 
 @register_microservice(name="s1", host="0.0.0.0", port=8083, endpoint="/v1/add")
@@ -58,19 +50,13 @@ def setUpClass(cls):
 
         cls.service_builder = ServiceOrchestrator()
 
-        cls.service_builder.add(cls.s0).add(cls.s1)
+        # simpler?: cls.service_builder.add(cls.s0).add(cls.s1)
+        cls.service_builder.add(opea_microservices["s0"]).add(opea_microservices["s1"])
         cls.service_builder.flow_to(cls.s0, cls.s1)
 
-        cls.service = MicroService(
-            "megaservice",
-            service_role=ServiceRoleType.MEGASERVICE,
-            host="0.0.0.0",
-            port=_MEGA_PORT,
-            endpoint="/whatever",
-            input_datatype=ChatCompletionRequest,
-            output_datatype=ChatCompletionResponse,
-        )
-        cls.service.start()
+        print(f"serving metrics from: http://localhost:{_METRIC_PORT}/metrics")
+        # requires prometheus_client >= 0.20.0 (earlier versions return None)
+        cls.server, cls.thread = start_http_server(_METRIC_PORT)
 
     @classmethod
     def tearDownClass(cls):
@@ -79,6 +65,9 @@ def tearDownClass(cls):
         cls.process1.terminate()
         cls.process2.terminate()
 
+        cls.server.shutdown()
+        cls.thread.join()
+
     async def test_schedule(self):
         result_dict, _ = await self.service_builder.schedule(initial_inputs={"text": "hello, "})
         response = result_dict["s1/MicroService"]
@@ -87,9 +76,10 @@ async def test_schedule(self):
         async for k in response.__reduce__()[2]["body_iterator"]:
             self.assertEqual(self.service_builder.extract_chunk_str(k).strip(), res_expected[idx])
             idx += 1
-        # check the metric increases from above
-        r = requests.get(f"http://localhost:{_MEGA_PORT}/metrics", timeout=10)
+        # check ServiceOrchestrator "megaservice_*" metrics resulting from above processing
+        r = await requests.get(f"http://localhost:{_METRIC_PORT}/metrics", timeout=5)
         self.assertEqual(r.status_code, 200)
+        # TODO: verify that the expected metrics are present
         print(r.text)
 
     def test_extract_chunk_str(self):