threaded output now added to tests during execution

Author: brandonjbjelland

Makefile

@@ -17,7 +17,7 @@ install: venv
 
 .PHONY: test
 test: install
-	$(COMMON_VARS) PROJECT_ID=$(PROJECT_ID) ${VENV_NAME}/bin/pytest -svvv --branch=$(SUBSTRATUS_BRANCH)
+	$(COMMON_VARS) PROJECT_ID=$(PROJECT_ID) ${VENV_NAME}/bin/pytest -s --branch=$(SUBSTRATUS_BRANCH)
 
 
 .PHONY: freeze

tests/capture_output_stream.py

@@ -3,13 +3,18 @@
 A utility to watch and stream outputs of local ipykernel events.
 """
 
-import asyncio
-from jupyter_client.asynchronous import AsyncKernelClient
+import logging
+from jupyter_client.blocking.client import BlockingKernelClient
 from watchdog.observers import Observer
 from watchdog.events import FileSystemEventHandler
 import glob
 import os
-import janus
+import re
+import threading
+from queue import Queue
+
+ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+logger = logging.getLogger(__name__)
 
 
 class NewFileHandler(FileSystemEventHandler):
@@ -19,8 +24,8 @@ def __init__(self, queue):
     def on_created(self, event):
         if event.is_directory or not event.src_path.endswith(".json"):
             return
-        print(f"New kernel detected: {event.src_path}")
-        self.queue.sync_q.put(event.src_path)
+        logger.debug(f"New kernel detected: {event.src_path}")
+        self.queue.put(event.src_path)
 
 
 def process_msg(msg):
@@ -30,52 +35,86 @@ def process_msg(msg):
             if "data" in msg["content"]
             else msg["content"].get("text", "")
         )
-        print(f"Output: {output}")
+        clean_text = ansi_escape.sub("", output)
+        logger.info(clean_text.rstrip("\n"))
     else:
-        print(f"Unhandled message type: {msg['msg_type']}")
-        print(msg)
+        logger.debug(f"Unhandled message type: {msg['msg_type']}")
+        logger.debug(msg)
 
 
-async def watch_kernel(connection_file):
-    kc = AsyncKernelClient(connection_file=connection_file)
+def watch_kernel(connection_file, stop_event):
+    kc = BlockingKernelClient(connection_file=connection_file)
     kc.load_connection_file()
     kc.start_channels()
 
-    while True:
-        msg = await kc.iopub_channel.get_msg()
-        process_msg(msg)
+    while not stop_event.is_set():
+        try:
+            msg = kc.get_iopub_msg(timeout=1)
+            if msg:
+                process_msg(msg)
+        except Exception as _:
+            continue
 
 
-async def watch_queue(queue, watched_files):
-    while True:
-        new_file = await queue.async_q.get()
+def watch_queue(queue, watched_files, stop_event):
+    while not stop_event.is_set():
+        new_file = queue.get()
         if new_file not in watched_files:
-            print(f"Processing new kernel: {new_file}")
+            logger.debug(f"Processing new kernel: {new_file}")
             watched_files.add(new_file)
-            asyncio.create_task(watch_kernel(new_file))
+            threading.Thread(target=watch_kernel, args=(new_file, stop_event)).start()
 
 
-async def main():
+def start_watches():
+    stop_event = threading.Event()
+    threads = []
     paths_to_watch = [
-        f"{os.path.expanduser('~')}/Library/Jupyter/runtime/",  # only works on mac OS
-        "/private/var/folders/9n/1rd9yjf913s10bzn5w9mdf_m0000gn/T/",  # I'm certain this is not portable as is
+        f"{os.path.expanduser('~')}/Library/Jupyter/runtime/",
+        "/private/var/folders/9n/1rd9yjf913s10bzn5w9mdf_m0000gn/T/",
         "/tmp/",
     ]
 
     existing_files = {
         f for path in paths_to_watch for f in glob.glob(os.path.join(path, "*.json"))
     }
-    print(f"Watching {len(existing_files)} existing files")
-
-    queue = janus.Queue()
-
-    tasks = [watch_kernel(config_file) for config_file in existing_files]
+    logger.info(f"Watching {len(existing_files)} existing files")
+
+    queue = Queue()
+    watched_files = set(existing_files)
+    for config_file in existing_files:
+        watch_thread = threading.Thread(
+            target=watch_kernel,
+            args=(
+                config_file,
+                stop_event,
+            ),
+        )
+        watch_thread.start()
+        threads.append(watch_thread)
 
     observer = Observer()
     for path in paths_to_watch:
         observer.schedule(NewFileHandler(queue), path, recursive=True)
 
-    observer_task = asyncio.get_event_loop().run_in_executor(None, observer.start)
-    queue_watcher = watch_queue(queue, existing_files)
+    observer.start()
+
+    # Start the watch_queue function in a separate thread
+    watch_queue_thread = threading.Thread(
+        target=watch_queue,
+        args=(
+            queue,
+            watched_files,
+            stop_event,
+        ),
+    )
+    watch_queue_thread.start()
+    threads.append(watch_queue_thread)
+    return threads, stop_event
+
+
+def main():
+    start_watches()
+
 
-    await asyncio.gather(*tasks, observer_task, queue_watcher)
+if __name__ == "__main__":
+    main()

tests/conftest.py

@@ -1,20 +1,16 @@
 import pytest
-import asyncio
 import os
 import subprocess
 
-from pytest_dependency import depends
 from testbook import testbook
 from google.cloud import storage
 import google.auth
 from google.auth.transport.requests import Request
 from testbook.testbook import TestbookNotebookClient
 import logging
-from capture_output_stream import main
+from capture_output_stream import start_watches
 
-logging.basicConfig(level=logging.INFO)
-# TODO(bjb): remove
-logging.basicConfig(level=logging.DEBUG)
+logging.basicConfig(level=logging.INFO, format="%(levelname)s:%(name)s:%(message)s")
 logger = logging.getLogger(__name__)
 
 
@@ -124,19 +120,40 @@ def auth_tb_serving_models(branch):
 
 @pytest.fixture(scope="session", autouse=True)
 def gcp_setup(auth_tb_quickstart):
-    logger.debug("Starting gcp_setup")
-
-    # Create a new event loop for this fixture
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-
-    logger.info("before capturing stream")
-    main_task = loop.create_task(main())  # Create task using the created loop
-    logger.info("output should stream")
-    # ... rest of the code ...
-
-    # Close the event loop when done
-    loop.close()
+    threads, stop_event = start_watches()  # start watches in threads
+    for attempt in range(3):  # Retry up to 3 times
+        logger.debug(f"Attempt {attempt} to execute installer gcp-up")
+        try:
+            auth_tb_quickstart.execute_cell("installer gcp-up")
+            assert "Apply complete!" in auth_tb_quickstart.cell_output_text(
+                "installer gcp-up"
+            )
+            break
+        except Exception as err:
+            logger.warning(f"gcp-up encountered an error: {err}")
+            if attempt == 1:
+                delete_state_lock()
+            continue
+
+    yield  # teardown below the yield
+
+    logger.debug("Tearing down gcp_setup")
+    for attempt in range(3):  # Retry up to 3 times
+        try:
+            auth_tb_quickstart.execute_cell("installer gcp-down")
+            assert "Apply complete!" in auth_tb_quickstart.cell_output_text(
+                "installer gcp-down"
+            )
+            break
+        except Exception as err:
+            logger.warning(f"gcp-down encountered an error: {err}")
+            if attempt == 1:
+                delete_state_lock()
+            continue
+        finally:
+            stop_event.set()
+            for thread in threads:
+                thread.join()  # Wait for threads to finish
 
 
 def delete_state_lock(