added working test

Author: Dylan Huang

tests/logging/test_elasticsearch_direct_http_handler.py

@@ -2,6 +2,7 @@
 import logging
 import time
 import pytest
+from datetime import datetime, timezone
 
 from eval_protocol.log_utils.elasticsearch_direct_http_handler import ElasticsearchDirectHttpHandler
 from eval_protocol.log_utils.elasticsearch_client import ElasticsearchClient
@@ -514,3 +515,91 @@ def test_elasticsearch_direct_http_handler_rollout_id_from_extra_overrides_env(
     )
 
     print(f"Successfully verified rollout_id override: extra '{extra_rollout_id}' overrode environment '{rollout_id}'")
+
+
+@pytest.mark.skipif(os.environ.get("CI") == "true", reason="Only run this test locally (skipped in CI)")
+def test_elasticsearch_direct_http_handler_timestamp_format(
+    elasticsearch_client: ElasticsearchClient, test_logger: logging.Logger, rollout_id: str
+):
+    """Test that ElasticsearchDirectHttpHandler formats timestamps correctly with UTC timezone."""
+
+    # Generate a unique test message
+    test_message = f"Timestamp format test message at {time.time()}"
+
+    # Record the time before logging to compare with the timestamp
+    before_log_time = datetime.now(timezone.utc)
+
+    # Send the log message
+    test_logger.info(test_message)
+
+    # Record the time after logging
+    after_log_time = datetime.now(timezone.utc)
+
+    # Give Elasticsearch a moment to process the document
+    time.sleep(3)
+
+    # Search for the document using the client
+    search_results = elasticsearch_client.search_by_match("message", test_message, size=1)
+
+    # Assert that we found our log message
+    assert search_results is not None, "Search should return results"
+    assert "hits" in search_results, "Search response should contain 'hits'"
+    assert "total" in search_results["hits"], "Search hits should contain 'total'"
+
+    total_hits = search_results["hits"]["total"]
+    if isinstance(total_hits, dict):
+        # Elasticsearch 7+ format
+        total_count = total_hits["value"]
+    else:
+        # Elasticsearch 6 format
+        total_count = total_hits
+
+    assert total_count > 0, f"Expected to find at least 1 log message, but found {total_count}"
+
+    # Verify the content of the found document
+    hits = search_results["hits"]["hits"]
+    assert len(hits) > 0, "Expected at least one hit"
+
+    found_document = hits[0]["_source"]
+    assert "@timestamp" in found_document, "Expected document to contain '@timestamp' field"
+
+    # Get the timestamp from the document
+    timestamp_str = found_document["@timestamp"]
+
+    # Verify the timestamp format matches ISO 8601 with UTC timezone (Z suffix)
+    assert timestamp_str.endswith("Z"), f"Expected timestamp to end with 'Z' (UTC), got: {timestamp_str}"
+
+    # Parse the timestamp to verify it's valid
+    try:
+        parsed_timestamp = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00"))
+    except ValueError as e:
+        pytest.fail(f"Failed to parse timestamp '{timestamp_str}': {e}")
+
+    # Verify the timestamp is timezone-aware (UTC)
+    assert parsed_timestamp.tzinfo is not None, "Expected timestamp to be timezone-aware"
+    utc_offset = parsed_timestamp.tzinfo.utcoffset(None)
+    assert utc_offset is not None and utc_offset.total_seconds() == 0, "Expected timestamp to be in UTC timezone"
+
+    # Verify the timestamp is within reasonable bounds (between before and after log time)
+    # Allow for some margin due to processing time
+    from datetime import timedelta
+
+    time_margin = timedelta(seconds=5)  # 5 seconds margin
+    assert before_log_time - time_margin <= parsed_timestamp <= after_log_time + time_margin, (
+        f"Expected timestamp {parsed_timestamp} to be between {before_log_time} and {after_log_time} "
+        f"(with {time_margin} margin)"
+    )
+
+    # Verify the timestamp format includes microseconds
+    assert "." in timestamp_str, "Expected timestamp to include microseconds"
+    assert timestamp_str.count(".") == 1, "Expected timestamp to have exactly one decimal point"
+
+    # Verify the format matches the expected pattern: YYYY-MM-DDTHH:MM:SS.ffffffZ
+    import re
+
+    iso_pattern = r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{6}Z$"
+    assert re.match(iso_pattern, timestamp_str), f"Expected timestamp to match ISO 8601 pattern, got: {timestamp_str}"
+
+    print(f"Successfully verified timestamp format: {timestamp_str}")
+    print(f"Parsed timestamp: {parsed_timestamp} (UTC)")
+    print(f"Timestamp is within expected time range: {before_log_time} <= {parsed_timestamp} <= {after_log_time}")