Add functional telemetry tests

hssyoo · hssyoo · commit ca2cba0f55a1 · 2025-05-20T12:43:16.000-04:00
diff --git a/.changes/next-release/enhancement-UserAgent-35014.json b/.changes/next-release/enhancement-UserAgent-35014.json
@@ -0,0 +1,5 @@
+{
+  "type": "enhancement",
+  "category": "User Agent",
+  "description": "Append session id to user agent string"
+}
diff --git a/awscli/telemetry.py b/awscli/telemetry.py
@@ -10,14 +10,14 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-import datetime
 import hashlib
 import io
 import os
 import socket
 import sqlite3
 import sys
 import threading
+import time
 from dataclasses import dataclass
 from functools import cached_property
 from pathlib import Path
@@ -51,8 +51,8 @@ class CLISessionDatabaseConnection:
     """
     _ENABLE_WAL = 'PRAGMA journal_mode=WAL'
 
-    def __init__(self):
-        self._connection = sqlite3.connect(
+    def __init__(self, connection=None):
+        self._connection = connection or sqlite3.connect(
             _CACHE_DIR / _DATABASE_FILENAME,
             check_same_thread=False,
             isolation_level=None,
@@ -165,19 +165,22 @@ def __init__(self, generator, writer, reader, sweeper):
 
         self._sweep_cache()
 
+    @cached_property
+    def cache_key(self):
+        return self._generator.generate_cache_key(self._hostname, self._tty)
+
     @cached_property
     def session_id(self):
-        cache_key = self._generator.generate_cache_key(
-            self._hostname, self._tty
-        )
-        if (cached_data := self._reader.read(cache_key)) is not None:
+        if (cached_data := self._reader.read(self.cache_key)) is not None:
             cached_data.timestamp = self._timestamp
             self._writer.write(cached_data)
             return cached_data.session_id
         session_id = self._generator.generate_session_id(
             self._hostname, self._tty, self._timestamp
         )
-        session_data = CLISessionData(cache_key, session_id, self._timestamp)
+        session_data = CLISessionData(
+            self.cache_key, session_id, self._timestamp
+        )
         self._writer.write(session_data)
         return session_id
 
@@ -200,7 +203,7 @@ def _hostname(self):
 
     @cached_property
     def _timestamp(self):
-        return int(datetime.datetime.now(datetime.timezone.utc).timestamp())
+        return int(time.time())
 
     def _sweep_cache(self):
         t = threading.Thread(
@@ -221,8 +224,8 @@ def _get_cli_session_orchestrator():
     )
 
 
-def add_session_id_component_to_user_agent_extra(session):
-    cli_session_orchestrator = _get_cli_session_orchestrator()
+def add_session_id_component_to_user_agent_extra(session, orchestrator):
+    cli_session_orchestrator = orchestrator or _get_cli_session_orchestrator()
     add_component_to_user_agent_extra(
         session, UserAgentComponent("sid", cli_session_orchestrator.session_id)
     )
diff --git a/tests/functional/test_telemetry.py b/tests/functional/test_telemetry.py
@@ -0,0 +1,256 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+import sqlite3
+from unittest.mock import MagicMock, patch
+
+import pytest
+from botocore.session import Session
+
+from awscli.telemetry import (
+    CLISessionData,
+    CLISessionDatabaseConnection,
+    CLISessionDatabaseReader,
+    CLISessionDatabaseSweeper,
+    CLISessionDatabaseWriter,
+    CLISessionGenerator,
+    CLISessionOrchestrator,
+    add_session_id_component_to_user_agent_extra,
+)
+from awscli.testutils import skip_if_windows
+
+
+@pytest.fixture
+def session_conn():
+    conn = CLISessionDatabaseConnection(
+        connection=sqlite3.connect(
+            # Use an in-memory db for testing.
+            ':memory:',
+            check_same_thread=False,
+            isolation_level=None,
+        ),
+    )
+    # Write an initial record.
+    conn.execute(
+        """
+            INSERT OR REPLACE INTO session (
+                key, session_id, timestamp
+            ) VALUES ('first_key', 'first_id', 5555555555)
+        """
+    )
+    return conn
+
+
+@pytest.fixture
+def session_writer(session_conn):
+    return CLISessionDatabaseWriter(session_conn)
+
+
+@pytest.fixture
+def session_reader(session_conn):
+    return CLISessionDatabaseReader(session_conn)
+
+
+@pytest.fixture
+def session_sweeper(session_conn):
+    return CLISessionDatabaseSweeper(session_conn)
+
+
+@pytest.fixture
+def session_generator():
+    return CLISessionGenerator()
+
+
+@pytest.fixture
+def expired_data(session_writer, session_reader, session_sweeper):
+    # Write an expired record.
+    session_writer.write(
+        CLISessionData(
+            key='expired_key',
+            session_id='expired_id',
+            timestamp=1000000000,
+        )
+    )
+    # Ensure expired record exists.
+    assert session_reader.read('expired_key') is not None
+    yield
+    # Ensure cleanup after test is run.
+    session_sweeper.sweep(1000000001)
+
+
+class TestCLISessionDatabaseConnection:
+    def test_ensure_database_setup(self, session_conn):
+        cursor = session_conn.execute(
+            """
+                SELECT name
+                FROM sqlite_master
+                WHERE type='table'
+                AND name='session';
+            """
+        )
+        assert cursor.fetchall() == [('session',)]
+
+    def test_timeout_does_not_raise_exception(self, session_conn):
+        class FakeConnection(sqlite3.Connection):
+            def execute(self, query, *parameters):
+                # Simulate timeout by always raising.
+                raise sqlite3.OperationalError()
+
+        fake_conn = CLISessionDatabaseConnection(FakeConnection(":memory:"))
+        cursor = fake_conn.execute(
+            """
+                SELECT name
+                FROM sqlite_master
+                WHERE type='table'
+                AND name='session';
+            """
+        )
+        assert cursor.fetchall() == []
+
+
+class TestCLISessionDatabaseWriter:
+    def test_write(self, session_writer, session_reader, session_sweeper):
+        session_writer.write(
+            CLISessionData(
+                key='new-key',
+                session_id='new-id',
+                timestamp=1000000000,
+            )
+        )
+        session_data = session_reader.read('new-key')
+        assert session_data.key == 'new-key'
+        assert session_data.session_id == 'new-id'
+        assert session_data.timestamp == 1000000000
+        session_sweeper.sweep(1000000001)
+
+
+class TestCLISessionDatabaseReader:
+    def test_read(self, session_reader):
+        session_data = session_reader.read('first_key')
+        assert session_data.key == 'first_key'
+        assert session_data.session_id == 'first_id'
+        assert session_data.timestamp == 5555555555
+
+    def test_read_nonexistent_record(self, session_reader):
+        session_data = session_reader.read('bad_key')
+        assert session_data is None
+
+
+class TestCLISessionDatabaseSweeper:
+    def test_sweep(self, expired_data, session_reader, session_sweeper):
+        session_sweeper.sweep(1000000001)
+        swept_data = session_reader.read('expired_key')
+        assert swept_data is None
+
+    def test_sweep_not_expired(
+        self, expired_data, session_reader, session_sweeper
+    ):
+        session_sweeper.sweep(1000000000)
+        swept_data = session_reader.read('expired_key')
+        assert swept_data is not None
+
+    def test_sweep_never_raises(self, session_sweeper):
+        # Normally this would raise `sqlite3.ProgrammingError`,
+        # but the `sweep` method catches bare exceptions.
+        session_sweeper.sweep({'bad': 'input'})
+
+
+class TestCLISessionGenerator:
+    def test_generate_session_id(self, session_generator):
+        session_id = session_generator.generate_session_id(
+            'my-hostname',
+            'my-tty',
+            1000000000,
+        )
+        assert session_id == 'd949713b13ee3fb52983b04316e8e6b5'
+
+    def test_generate_cache_key(self, session_generator):
+        cache_key = session_generator.generate_cache_key(
+            'my-hostname',
+            'my-tty',
+        )
+        assert cache_key == 'b1ca2be0ffac12f172933b6777e06f2c'
+
+
+@skip_if_windows("No os.ttyname")
+@patch('sys.stdin')
+@patch('time.time', return_value=5555555555)
+@patch('socket.gethostname', return_value='my-hostname')
+@patch('os.ttyname', return_value='my-tty')
+class TestCLISessionOrchestrator:
+    def test_session_id_gets_cached(
+        self,
+        patched_tty_name,
+        patched_hostname,
+        patched_time,
+        patched_stdin,
+        session_sweeper,
+        session_generator,
+        session_reader,
+        session_writer,
+    ):
+        patched_stdin.fileno.return_value = None
+        orchestrator = CLISessionOrchestrator(
+            session_generator, session_writer, session_reader, session_sweeper
+        )
+        assert orchestrator.session_id == '881cea8546fa4888970cce8d133c3bf9'
+
+        session_data = session_reader.read(orchestrator.cache_key)
+        assert session_data.key == orchestrator.cache_key
+        assert session_data.session_id == orchestrator.session_id
+        assert session_data.timestamp == 5555555555
+
+    def test_cached_session_id_gets_updated(
+        self,
+        patched_tty_name,
+        patched_hostname,
+        patched_time,
+        patched_stdin,
+        session_sweeper,
+        session_generator,
+        session_reader,
+        session_writer,
+    ):
+        patched_stdin.fileno.return_value = None
+
+        # First, generate and cache a session id.
+        orchestrator_1 = CLISessionOrchestrator(
+            session_generator, session_writer, session_reader, session_sweeper
+        )
+        session_id_1 = orchestrator_1.session_id
+        session_data_1 = session_reader.read(orchestrator_1.cache_key)
+        assert session_data_1.session_id == session_id_1
+
+        # Update the timestamp and get the new session id.
+        patched_time.return_value = 7777777777
+        orchestrator_2 = CLISessionOrchestrator(
+            session_generator, session_writer, session_reader, session_sweeper
+        )
+        session_id_2 = orchestrator_2.session_id
+        session_data_2 = session_reader.read(orchestrator_2.cache_key)
+
+        # Cache key should be the same.
+        assert session_data_2.key == session_data_1.key
+        # Session id and timestamp should be updated.
+        assert session_data_2.session_id == session_id_2
+        assert session_data_2.session_id != session_data_1.session_id
+        assert session_data_2.timestamp == 7777777777
+        assert session_data_2.timestamp != session_data_1.timestamp
+
+
+def test_add_session_id_component_to_user_agent_extra():
+    session = MagicMock(Session)
+    session.user_agent_extra = ''
+    orchestrator = MagicMock(CLISessionOrchestrator)
+    orchestrator.session_id = 'my-session-id'
+    add_session_id_component_to_user_agent_extra(session, orchestrator)
+    assert session.user_agent_extra == 'sid/my-session-id'
