add audio test

src/nv_ingest/framework/orchestration/morpheus/util/pipeline/stage_builders.py

@@ -408,29 +408,38 @@ def get_audio_retrieval_service(env_var_prefix):
         "",
     )
     auth_token = os.environ.get(
-        "RIVA_NGC_API_KEY",
+        "NVIDIA_BUILD_API_KEY",
+        "",
+    ) or os.environ.get(
+        "NGC_API_KEY",
         "",
     )
     infer_protocol = os.environ.get(
         "AUDIO_INFER_PROTOCOL",
         "http" if http_endpoint else "grpc" if grpc_endpoint else "",
     )
+    function_id = os.environ.get(
+        "AUDIO_FUNCTION_ID",
+        "",
+    )
 
-    logger.info(f"{prefix}_GRPC_TRITON: {grpc_endpoint}")
-    logger.info(f"{prefix}_HTTP_TRITON: {http_endpoint}")
+    logger.info(f"{prefix}_GRPC_ENDPOINT: {grpc_endpoint}")
+    logger.info(f"{prefix}_HTTP_ENDPOINT: {http_endpoint}")
     logger.info(f"{prefix}_INFER_PROTOCOL: {infer_protocol}")
+    logger.info(f"{prefix}_FUNCTION_ID: {function_id}")
 
-    return grpc_endpoint, http_endpoint, auth_token, infer_protocol
+    return grpc_endpoint, http_endpoint, auth_token, infer_protocol, function_id
 
 
 def add_audio_extractor_stage(pipe, morpheus_pipeline_config, ingest_config, default_cpu_count):
-    audio_grpc, audio_http, audio_auth, audio_infer_protocol = get_audio_retrieval_service("audio")
+    audio_grpc, audio_http, audio_auth, audio_infer_protocol, audio_function_id = get_audio_retrieval_service("audio")
     audio_extractor_config = ingest_config.get(
         "audio_extraction_module",
         {
             "audio_extraction_config": {
                 "audio_endpoints": (audio_grpc, audio_http),
                 "audio_infer_protocol": audio_infer_protocol,
+                "function_id": audio_function_id,
                 "auth_token": audio_auth,
                 # All auth tokens are the same for the moment
             }

tests/integration/test_extract_audio.py

@@ -0,0 +1,32 @@
+import subprocess
+import sys
+import time
+
+import pytest
+from nv_ingest_api.util.message_brokers.simple_message_broker import SimpleClient
+from nv_ingest_client.client import Ingestor
+from nv_ingest_client.client import NvIngestClient
+
+
+def test_images_extract_only(
+    pipeline_process,
+):
+    client = NvIngestClient(
+        message_client_allocator=SimpleClient,
+        message_client_port=7671,
+        message_client_hostname="localhost",
+    )
+
+    ingestor = (
+        Ingestor(client=client)
+        .files("./data/multimodal_test.wav")
+        .extract(
+            extract_method="audio",
+        )
+    )
+
+    results = ingestor.ingest()
+
+    assert len(results) == 1
+
+    print(results)