From c2b2082cc45aa12d934b5f2309f1656144bc81fa Mon Sep 17 00:00:00 2001 From: Jeremy Prevost Date: Thu, 19 Mar 2026 15:38:36 -0400 Subject: [PATCH 1/4] Removes example lambda function and test files Why are these changes being introduced: * the example lambda function and test files are no longer needed and are being removed to clean up the codebase --- lambdas/my_function.py | 25 ------------------------- tests/test_my_function.py | 7 ------- 2 files changed, 32 deletions(-) delete mode 100644 lambdas/my_function.py delete mode 100644 tests/test_my_function.py diff --git a/lambdas/my_function.py b/lambdas/my_function.py deleted file mode 100644 index 09b0ba8..0000000 --- a/lambdas/my_function.py +++ /dev/null @@ -1,25 +0,0 @@ -import json -import logging - -from lambdas.config import Config, configure_logger, configure_sentry - -# --------------------------------------- -# One-time, Lambda cold start setup -# --------------------------------------- -CONFIG = Config() - -root_logger = logging.getLogger() -log_config_message = configure_logger(root_logger) -logger = logging.getLogger(__name__) -logger.info(log_config_message) - -configure_sentry() - - -# --------------------------------------- -# Lambda handler entrypoint -# --------------------------------------- -def lambda_handler(event: dict, lambda_context: dict) -> str: - logger.debug(json.dumps(event)) - logger.info("LaLambda context: %s", lambda_context) - return "You have successfully called this lambda!" diff --git a/tests/test_my_function.py b/tests/test_my_function.py deleted file mode 100644 index d8b187a..0000000 --- a/tests/test_my_function.py +++ /dev/null @@ -1,7 +0,0 @@ -from lambdas import my_function - - -def test_my_function(): - assert ( - my_function.lambda_handler({}, {}) == "You have successfully called this lambda!" - ) From 81aeb23991b22a99ebaa21833b50d30097a95792 Mon Sep 17 00:00:00 2001 From: Jeremy Prevost Date: Thu, 19 Mar 2026 16:09:06 -0400 Subject: [PATCH 2/4] Adds ping support Why are these changes being introduced: * Adds a health check/ping endpoint to the tokenizer lambda, which can be used by AWS to determine if the lambda is healthy and ready to receive traffic. * This is important for ensuring high availability and reliability of the service. * This also allows us to keep the lambda warm, which can improve performance by reducing cold start latency. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/USE-431 How does this address that need: * Adds check for "ping" key in the event, and if present, returns a simple response indicating the service is healthy. Document any side effects to this change: * We cannot use the "ping" key in the event for any other purpose, as it is now reserved for health checks. --- lambdas/tokenizer_handler.py | 7 +++++++ tests/test_tokenizer_handler.py | 17 +++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/lambdas/tokenizer_handler.py b/lambdas/tokenizer_handler.py index 051bc1f..615c51a 100644 --- a/lambdas/tokenizer_handler.py +++ b/lambdas/tokenizer_handler.py @@ -43,6 +43,13 @@ def lambda_handler(event: dict, lambda_context: Context) -> dict: query_tokenizer = _get_tokenizer() + # Handle ping or health check events + # We add this after query_tokenizer initialization to ensure the tokenizer is + # initialized during cold start, even for pings + if event.get("ping"): + logger.info("Received ping event") + return {"status": "ok"} + # Generate query tokens query = event.get("query", "") diff --git a/tests/test_tokenizer_handler.py b/tests/test_tokenizer_handler.py index 9c21cba..7ab45b3 100644 --- a/tests/test_tokenizer_handler.py +++ b/tests/test_tokenizer_handler.py @@ -59,3 +59,20 @@ def test_returns_error_for_whitespace_only_query(mock_query_tokenizer): result = tokenizer_handler.lambda_handler({"query": " "}, {}) assert result == {"error": "Query is required in the event payload."} mock_query_tokenizer.tokenize_query.assert_not_called() + + +def test_ping_event_returns_ok_status(mock_query_tokenizer): + result = tokenizer_handler.lambda_handler({"ping": True}, {}) + assert result == {"status": "ok"} + + +def test_ping_event_does_not_call_tokenizer(mock_query_tokenizer): + tokenizer_handler.lambda_handler({"ping": True}, {}) + mock_query_tokenizer.tokenize_query.assert_not_called() + + +def test_non_ping_event_is_not_short_circuited(mock_query_tokenizer): + mock_query_tokenizer.tokenize_query.return_value = {"hello": 1.5} + result = tokenizer_handler.lambda_handler({"query": "hello"}, {}) + assert "query" in result + mock_query_tokenizer.tokenize_query.assert_called_once() From cbf00025447b031b4819cbd12827d1999f7222af Mon Sep 17 00:00:00 2001 From: Jeremy Prevost Date: Thu, 19 Mar 2026 16:09:22 -0400 Subject: [PATCH 3/4] Adds integration tests These tests are not strictly necessary, but I was nervous we never actually tested the tokenizer handler with a real tokenizer. These tests load a real tokenizer and test that the handler can generate tokens for a simple query. They are marked as "integration" tests since they load real models from disk and are slower than typical unit tests. They are enabled by default, but can be skipped with `pytest -m "not integration"`. The total test time does not increase enough to make the default be to skip these tests at this time, but if we need to in the future we can change the default. --- pyproject.toml | 3 +++ tests/test_query_tokenizer.py | 36 +++++++++++++++++++++++++++++++++ tests/test_tokenizer_handler.py | 15 ++++++++++++++ 3 files changed, 54 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 8a0c82f..c5a0e7e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,9 @@ exclude = ["tests/"] [tool.pytest.ini_options] log_level = "INFO" +markers = [ + "integration: tests that load real models/files from disk (slow)", +] [tool.ruff] target-version = "py314" diff --git a/tests/test_query_tokenizer.py b/tests/test_query_tokenizer.py index e1d450d..41de386 100644 --- a/tests/test_query_tokenizer.py +++ b/tests/test_query_tokenizer.py @@ -129,3 +129,39 @@ def test_load_idf_raises_filenotfounderror_when_idf_file_missing(): pytest.raises(FileNotFoundError), ): QueryTokenizer() + + +# --------------------------------------------------------------------------- +# Integration tests — load the real tokenizer and IDF from disk +# Run only integration tests: uv run pytest -m integration +# Run all except integration tests: uv run pytest -m "not integration" +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_integration_tokenize_query_returns_nonempty_dict(): + """Real tokenizer and IDF: tokenize_query returns a non-empty dict.""" + qt = QueryTokenizer() + result = qt.tokenize_query("machine learning") + assert isinstance(result, dict) + assert len(result) > 0 + + +@pytest.mark.integration +def test_integration_tokenize_query_values_are_positive_floats(): + """Real tokenizer and IDF: all weights are positive floats.""" + qt = QueryTokenizer() + result = qt.tokenize_query("machine learning") + for token, weight in result.items(): + assert isinstance(token, str), f"Expected str key, got {type(token)}" + assert isinstance(weight, float), f"Expected float weight, got {type(weight)}" + assert weight > 0, f"Expected positive weight for {token!r}, got {weight}" + + +@pytest.mark.integration +def test_integration_tokenize_query_deterministic(): + """Real tokenizer and IDF: same input always produces the same output.""" + qt = QueryTokenizer() + result_a = qt.tokenize_query("open access repositories") + result_b = qt.tokenize_query("open access repositories") + assert result_a == result_b diff --git a/tests/test_tokenizer_handler.py b/tests/test_tokenizer_handler.py index 7ab45b3..30f4f67 100644 --- a/tests/test_tokenizer_handler.py +++ b/tests/test_tokenizer_handler.py @@ -76,3 +76,18 @@ def test_non_ping_event_is_not_short_circuited(mock_query_tokenizer): result = tokenizer_handler.lambda_handler({"query": "hello"}, {}) assert "query" in result mock_query_tokenizer.tokenize_query.assert_called_once() + + +# --------------------------------------------------------------------------- +# Integration tests — load the real tokenizer and IDF from disk +# Run only integration tests: uv run pytest -m integration +# Run all except integration tests: uv run pytest -m "not integration" +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_integration_lambda_handler_returns_opensearch_query(): + tokenizer_handler._get_tokenizer.cache_clear() # ensure cold start + result = tokenizer_handler.lambda_handler({"query": "open access"}, {}) + assert "query" in result + assert "bool" in result["query"] From 1d5091cffcd5d635d461acea776c3d0a3c7c42f5 Mon Sep 17 00:00:00 2001 From: Jeremy Prevost Date: Fri, 20 Mar 2026 08:28:30 -0400 Subject: [PATCH 4/4] Address automated feedback --- lambdas/tokenizer_handler.py | 4 ++-- tests/test_tokenizer_handler.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/lambdas/tokenizer_handler.py b/lambdas/tokenizer_handler.py index 615c51a..e1aa5e1 100644 --- a/lambdas/tokenizer_handler.py +++ b/lambdas/tokenizer_handler.py @@ -46,8 +46,8 @@ def lambda_handler(event: dict, lambda_context: Context) -> dict: # Handle ping or health check events # We add this after query_tokenizer initialization to ensure the tokenizer is # initialized during cold start, even for pings - if event.get("ping"): - logger.info("Received ping event") + if "ping" in event: + logger.debug("Received ping event") return {"status": "ok"} # Generate query tokens diff --git a/tests/test_tokenizer_handler.py b/tests/test_tokenizer_handler.py index 30f4f67..a7624fc 100644 --- a/tests/test_tokenizer_handler.py +++ b/tests/test_tokenizer_handler.py @@ -66,7 +66,7 @@ def test_ping_event_returns_ok_status(mock_query_tokenizer): assert result == {"status": "ok"} -def test_ping_event_does_not_call_tokenizer(mock_query_tokenizer): +def test_ping_event_does_not_call_tokenize_query(mock_query_tokenizer): tokenizer_handler.lambda_handler({"ping": True}, {}) mock_query_tokenizer.tokenize_query.assert_not_called()