test_user_agent.py

from __future__ import absolute_import

import json
import os
from mock import patch, mock_open

import pytest

from sagemaker.core.telemetry.attribution import _CREATED_BY_ENV_VAR
from sagemaker.core.utils.user_agent import (
    SagemakerCore_PREFIX,
    SagemakerCore_VERSION,
    NOTEBOOK_PREFIX,
    STUDIO_PREFIX,
    process_notebook_metadata_file,
    process_studio_metadata_file,
    get_user_agent_extra_suffix,
    sanitize_user_agent_string_component,
)


@pytest.fixture(autouse=True)
def clean_env():
    yield
    if _CREATED_BY_ENV_VAR in os.environ:
        del os.environ[_CREATED_BY_ENV_VAR]


# Test process_notebook_metadata_file function
def test_process_notebook_metadata_file_exists(tmp_path):
    notebook_file = tmp_path / "sagemaker-notebook-instance-version.txt"
    notebook_file.write_text("instance_type")

    with patch("os.path.exists", return_value=True):
        with patch("builtins.open", mock_open(read_data=notebook_file.read_text())):
            assert process_notebook_metadata_file() == "instance_type"


def test_process_notebook_metadata_file_not_exists(tmp_path):
    with patch("os.path.exists", return_value=False):
        assert process_notebook_metadata_file() is None


# Test process_studio_metadata_file function
def test_process_studio_metadata_file_exists(tmp_path):
    studio_file = tmp_path / "resource-metadata.json"
    studio_file.write_text(json.dumps({"AppType": "studio_type"}))

    with patch("os.path.exists", return_value=True):
        with patch("builtins.open", mock_open(read_data=studio_file.read_text())):
            assert process_studio_metadata_file() == "studio_type"


def test_process_studio_metadata_file_not_exists(tmp_path):
    with patch("os.path.exists", return_value=False):
        assert process_studio_metadata_file() is None


# Test sanitize_user_agent_string_component function
def test_sanitize_replaces_slash_with_dash():
    assert sanitize_user_agent_string_component("awslabs/agent-plugins/sagemaker-ai") == "awslabs-agent-plugins-sagemaker-ai"


def test_sanitize_allows_alphanumeric():
    assert sanitize_user_agent_string_component("abc123") == "abc123"


def test_sanitize_replaces_hash_when_not_allowed():
    assert sanitize_user_agent_string_component("foo#bar") == "foo-bar"


def test_sanitize_allows_hash_when_permitted():
    assert sanitize_user_agent_string_component("foo#bar", allow_hash=True) == "foo#bar"


def test_sanitize_replaces_space_with_dash():
    assert sanitize_user_agent_string_component("foo bar") == "foo-bar"


# Test get_user_agent_extra_suffix function
def test_get_user_agent_extra_suffix():
    assert get_user_agent_extra_suffix() == f"lib/{SagemakerCore_PREFIX}#{SagemakerCore_VERSION}"

    with patch(
        "sagemaker.core.utils.user_agent.process_notebook_metadata_file",
        return_value="instance_type",
    ):
        assert (
            get_user_agent_extra_suffix()
            == f"lib/{SagemakerCore_PREFIX}#{SagemakerCore_VERSION} md/{NOTEBOOK_PREFIX}#instance_type"
        )

    with patch(
        "sagemaker.core.utils.user_agent.process_studio_metadata_file", return_value="studio_type"
    ):
        assert (
            get_user_agent_extra_suffix()
            == f"lib/{SagemakerCore_PREFIX}#{SagemakerCore_VERSION} md/{STUDIO_PREFIX}#studio_type"
        )


def test_get_user_agent_extra_suffix_without_created_by():
    suffix = get_user_agent_extra_suffix()
    assert "createdBy" not in suffix


def test_get_user_agent_extra_suffix_with_created_by():
    os.environ[_CREATED_BY_ENV_VAR] = "awslabs/agent-plugins/sagemaker-ai"
    suffix = get_user_agent_extra_suffix()
    assert "md/createdBy#awslabs-agent-plugins-sagemaker-ai" in suffix


def test_get_user_agent_extra_suffix_created_by_sanitized():
    os.environ[_CREATED_BY_ENV_VAR] = "my agent/v1.0 (test)"
    suffix = get_user_agent_extra_suffix()
    assert "md/createdBy#my-agent-v1.0--test-" in suffix