chore(http-client): full test-coverage 1/2

[viraatc]

What does this PR do?

with this MR:

• reworked cached-headers api to be simpler (no side effects).
• reached 100% coverage on http.py

src/inference_endpoint/endpoint_client/__init__.py                                3      0   100%
src/inference_endpoint/endpoint_client/accumulator_protocol.py                    3      0   100%
src/inference_endpoint/endpoint_client/adapter_protocol.py                       18      0   100%
src/inference_endpoint/endpoint_client/config.py                                 74      4    95%   142, 177, 199, 203
src/inference_endpoint/endpoint_client/cpu_affinity.py                          214     50    77%   51, 96, 101, 111-124, 156-157, 171, 173, 203, 221, 248-254, 304-305, 314-316, 340-342, 358-359, 373, 380, 394-396, 422-435, 505, 507, 525
src/inference_endpoint/endpoint_client/http.py                                  342      0   100%
src/inference_endpoint/endpoint_client/http_client.py                            64      3    95%   69, 86, 158
src/inference_endpoint/endpoint_client/http_sample_issuer.py                     38      8    79%   72-79, 84-86
src/inference_endpoint/endpoint_client/utils.py                                  49     15    69%   40-59, 70, 72, 85, 91-92
src/inference_endpoint/endpoint_client/worker.py                                246     60    76%   87-88, 102-103, 121-123, 160, 193, 223-255, 266-280, 284-286, 323, 342-347, 384-385, 413-419, 437, 467, 480, 505, 516, 526, 560-561, 572-575
src/inference_endpoint/endpoint_client/worker_manager.py                         77      7    91%   98-99, 105, 146, 154, 163-164

Type of change

• Bug fix
• New feature
• Documentation update
• Refactor/cleanup

Related issues

Testing

• Tests added/updated
• All tests pass locally
• Manual testing completed

Checklist

• Code follows project style
• Pre-commit hooks pass
• Documentation updated (if needed)

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces significant improvements to the core HTTP client by refactoring its header caching mechanism for better performance and predictability. Concurrently, it vastly expands the unit test suite, ensuring comprehensive coverage of the HTTP protocol handling and connection pooling logic. These changes aim to make the HTTP client more robust, efficient, and easier to maintain by addressing edge cases and optimizing common request patterns.

Highlights

• Refactored HTTP Request Header Caching: The HttpRequestTemplate class was updated to simplify header caching. It now uses a single cached_headers attribute for pre-encoded static headers, moving away from a dynamic dictionary-based cache for extra_headers. This optimizes for scenarios where a fixed set of headers is frequently sent.
• Enhanced HTTP Client Test Coverage: Extensive new unit tests were added for HttpResponseProtocol and ConnectionPool. These tests cover various aspects including connection lifecycle, error propagation, idle timeouts, stale connection handling, pool warmup, and fast-path optimizations, significantly improving the robustness of the HTTP client.
• Optimized 'Slow Path' for Dynamic Headers: The build_request method's 'slow path' for extra_headers now directly encodes headers per-call, with the expectation that frequently used headers are pre-encoded via cache_headers() during setup, streamlining the request building process.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• src/inference_endpoint/endpoint_client/http.py
  • Simplified HttpRequestTemplate's __slots__ to use cached_headers instead of _extra_headers_cache and extra_cached_headers.
  • Updated __init__ to initialize cached_headers directly.
  • Modified cache_headers to pre-encode headers directly into self.cached_headers rather than using a dynamic cache.
  • Adjusted build_request to remove dynamic caching for extra_headers, making it a per-call encoding path, and updated comments to reflect this change.
• tests/unit/endpoint_client/test_http.py
  • Removed redundant comments across multiple tests for cleaner code.
  • Updated test_build_request_with_extra_headers docstring to specify 'slow path' testing.
  • Removed test_extra_headers_are_cached due to the refactored header caching mechanism.
  • Updated test_cache_headers_pre_caches to assert against the new cached_headers attribute.
  • Renamed test_build_request_with_cache_headers_pre_caches to test_build_request_with_cached_headers and adapted it to the new caching logic.
  • Renamed test_build_request_without_extra_headers to test_build_request_fast_path and added an assertion for the Accept header.
  • Added test_streaming_accept_header to verify the Accept header for streaming requests.
  • Added an assertion for protocol.transport in test_parses_200_response_with_body.
  • Reduced asyncio.sleep durations to 0 in test_iter_body_yields_chunks.
  • Added test_iter_body_pre_buffered_and_already_complete to test pre-buffered chunk consumption.
  • Added test_connection_lost_before_body_complete as a parameterized test for connection loss during body reception.
  • Added test_connection_lost_with_complete_message to verify behavior when connection is lost after a complete message.
  • Added test_data_received_parser_error as a parameterized test for HttpParserError propagation.
  • Added test_read_body_and_headers_fast_paths to test immediate returns for already parsed headers/bodies.
  • Updated test_close_cleans_up_all to close both idle and in-use connections.
  • Added test_stale_connection_discarded_on_acquire to ensure dead connections are discarded.
  • Added test_idle_timeout_discards_connection to test connection discarding based on idle time.
  • Added test_warmup as a parameterized test for pre-establishing connections in the pool.
  • Added test_close_cancels_waiters to verify that pool closure cancels waiting tasks.
  • Added test_release_idempotent to confirm that releasing a connection multiple times has no adverse effect.
  • Added test_unlimited_pool to test connection pools configured for unlimited connections.
  • Added test_waiter_creates_new_connection to test new connection creation when waiters are unblocked.
  • Added test_is_stale_various_conditions to thoroughly test the is_stale logic under different conditions.

Activity

• The author, viraatc, initiated this pull request as a new feature.
• Tests have been added and updated to cover the new functionality and existing components.
• All local tests are reported to be passing.
• Manual testing has been completed by the author.
• The code adheres to project style guidelines.
• Pre-commit hooks have passed successfully.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[Copilot]

Pull request overview

This PR adds comprehensive unit test coverage for the HTTP API (HttpResponseProtocol, ConnectionPool, and HttpRequestTemplate) and simplifies the HttpRequestTemplate slow path by removing the per-call extra_headers dict-cache in favour of a single cached_headers bytes buffer populated once at setup time via cache_headers().

Changes:

• Removes the _extra_headers_cache dict from HttpRequestTemplate; the slow path now re-encodes extra_headers on every call, and callers are expected to use cache_headers() for headers that repeat on every request.
• Renames extra_cached_headers → cached_headers (now public).
• Adds many new unit tests covering HttpResponseProtocol error paths (connection_lost, eof_received, data_received parser errors, iter_body pre-buffered case, read_body/read_headers fast paths) and ConnectionPool edge cases (stale-connection discard, idle timeout, waiter cancellation, unlimited pool, warmup, idempotent release).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
src/inference_endpoint/endpoint_client/http.py	Removes _extra_headers_cache, renames extra_cached_headers → cached_headers, simplifies cache_headers() and slow-path in build_request()
tests/unit/endpoint_client/test_http.py	Deletes now-invalid cache tests, updates existing tests to match the new API, and adds ~10 new test cases for previously uncovered code paths

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[gemini-code-assist]

Code Review

The pull request refactors the HttpRequestTemplate class to simplify header caching. It replaces a dictionary-based cache (_extra_headers_cache) with a single byte string (cached_headers) for pre-encoded headers. The cache_headers method now directly appends new headers to this byte string, removing the logic for preventing duplicate header sets. The build_request method's 'slow path' for extra_headers no longer caches these headers, relying solely on cache_headers for pre-encoding. This change introduces a high-severity HTTP Header Injection vulnerability in both the cache_headers and build_request methods, as header values are directly embedded into strings without sanitization, allowing attackers to inject newline characters and arbitrary headers. Additionally, the simplified cache_headers method loses protection against duplicate headers, potentially leading to malformed requests if called multiple times with the same headers. The associated unit tests were updated to reflect these changes, including removing tests for the old caching mechanism and adding new tests for the updated HttpResponseProtocol behavior, such as handling connection loss and parser errors.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[arekay-nv]

LGTM!
Thanks.