Websocket Integration Test (#907)

* Websocket integration test

containers/orchestration/app/config.py

@@ -1,6 +1,6 @@
 from functools import lru_cache
-from pydantic import BaseSettings
 from typing import Optional
+from pydantic import BaseSettings
 
 
 class Settings(BaseSettings):

containers/orchestration/app/main.py

@@ -9,6 +9,7 @@
     File,
     HTTPException,
     WebSocket,
+    WebSocketDisconnect,
 )
 from typing import Annotated, Optional
 from pathlib import Path
@@ -72,21 +73,26 @@ async def process_message_endpoint_ws(
     """
 
     await websocket.accept()
-    while True:
-        file = await websocket.receive_bytes()
+    try:
+        while True:
+            file = await websocket.receive_bytes()
 
-        zipped_file = ZipFile(io.BytesIO(file), "r")
-        unzipped_file = unzip_ws(zipped_file)
+            zipped_file = ZipFile(io.BytesIO(file), "r")
+            unzipped_file = unzip_ws(zipped_file)
 
-        # Hardcoded message_type for MVP
-        input = {
-            "message_type": "eicr",
-            "include_error_types": "errors",
-            "message": unzipped_file,
-        }
+            # Hardcoded message_type for MVP
+            input = {
+                "message_type": "eicr",
+                "include_error_types": "errors",
+                "message": unzipped_file,
+            }
 
-        processing_config = load_processing_config("sample-orchestration-config.json")
-        await call_apis(config=processing_config, input=input, websocket=websocket)
+            processing_config = load_processing_config(
+                "sample-orchestration-config.json"
+            )
+            await call_apis(config=processing_config, input=input, websocket=websocket)
+    except WebSocketDisconnect:
+        await websocket.close()
 
 
 @app.post("/process", status_code=200, responses=process_message_response_examples)

containers/orchestration/app/services.py

@@ -85,12 +85,14 @@ async def call_apis(
         service = step["service"]
         endpoint = step["endpoint"]
         f = f"{service}_payload"
+
         if f in globals() and callable(globals()[f]) and service_urls[service]:
             function_to_call = globals()[f]
             payload = function_to_call(
                 input=input, response=response, step=step, config=config
             )
             url = service_urls[service] + step["endpoint"]
+            url = url.replace('"', "")
             print(f"Url: {url}")
             response = post_request(url, payload)
             print(f"Status Code: {response.status_code}")
@@ -111,4 +113,4 @@ async def call_apis(
                 detail="The Building Block you are attempting to call does not exist:"
                 + f" {service}",
             )
-    return (response, responses)
+    return response, responses

containers/orchestration/local-dev.env

@@ -1,7 +1,7 @@
-VALIDATION_URL="http://localhost:8081"
-FHIR_CONVERTER_URL="http://localhost:8082"
-INGESTION_URL="http://localhost:8083"
-MESSAGE_PARSER_URL="http://localhost:8085"
-SMARTY_AUTH_ID="placeholder"
-SMARTY_AUTH_TOKEN="placeholder"
-LICENSE_TYPE="us-rooftop-geocoding-enterprise-cloud"
+VALIDATION_URL=http://localhost:8081
+FHIR_CONVERTER_URL=http://localhost:8082
+INGESTION_URL=http://localhost:8083
+MESSAGE_PARSER_URL=http://localhost:8085
+SMARTY_AUTH_ID=placeholder
+SMARTY_AUTH_TOKEN=placeholder
+LICENSE_TYPE=us-rooftop-geocoding-enterprise-cloud

containers/orchestration/tests/integration/test.env

@@ -1,7 +1,7 @@
-VALIDATION_URL="http://validation-service:8080"
-FHIR_CONVERTER_URL="http://fhir-converter-service:8080"
-INGESTION_URL="http://ingestion-service:8080"
-MESSAGE_PARSER_URL="http://message-parser-service:8080"
+VALIDATION_URL=http://validation-service:8080
+FHIR_CONVERTER_URL=http://fhir-converter-service:8080
+INGESTION_URL=http://ingestion-service:8080
+MESSAGE_PARSER_URL=http://message-parser-service:8080
 SMARTY_AUTH_ID="placeholder"
 SMARTY_AUTH_TOKEN="placeholder"
 LICENSE_TYPE="us-rooftop-geocoding-enterprise-cloud"

containers/orchestration/tests/integration/test_orchestration.py

@@ -2,7 +2,11 @@
 import pytest
 import os
 from pathlib import Path
+from app.config import get_settings
+from app.main import app
+from starlette.testclient import TestClient
 
+get_settings()
 
 ORCHESTRATION_URL = "http://localhost:8080"
 PROCESS_ENDPOINT = ORCHESTRATION_URL + "/process"
@@ -69,3 +73,51 @@ def test_process_endpoint_with_zip(setup):
         )
         assert orchestration_response.status_code == 200
         assert orchestration_response.json()["message"] == "Processing succeeded!"
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_websocket_process_message_endpoint():
+    expected_response_message = {
+        "steps": [
+            {"endpoint": "/validate", "service": "validation"},
+            {"endpoint": "/convert-to-fhir", "service": "fhir_converter"},
+            {
+                "endpoint": "/fhir/harmonization/standardization/standardize_names",
+                "service": "ingestion",
+            },
+            {
+                "endpoint": "/fhir/harmonization/standardization/standardize_phones",
+                "service": "ingestion",
+            },
+            {
+                "endpoint": "/fhir/harmonization/standardization/standardize_dob",
+                "service": "ingestion",
+            },
+            {"endpoint": "/parse_message", "service": "message_parser"},
+        ],
+        "validate": {"Message": "OK", "status_code": 200},
+    }
+
+    client = TestClient(app)
+
+    # Pull in and read test zip file
+    with open(
+        Path(__file__).parent.parent.parent.parent.parent
+        / "tests"
+        / "assets"
+        / "orchestration"
+        / "test_zip.zip",
+        "rb",
+    ) as file:
+        test_zip = file.read()
+
+    # Create fake websocket connection
+    with client.websocket_connect("/process-ws") as websocket:
+        # Send zip into fake connection, triggering process-ws endpoint
+        websocket.send_bytes(test_zip)
+
+        # Pull response message from websocket connection like frontend would
+        messages = websocket.receive_json()
+
+    assert messages == expected_response_message