diff --git a/src/xpk/core/config.py b/src/xpk/core/config.py index 44055075d..773354cac 100644 --- a/src/xpk/core/config.py +++ b/src/xpk/core/config.py @@ -31,6 +31,7 @@ CLUSTER_NAME_KEY = 'cluster-name' PROJECT_KEY = 'project-id' CLIENT_ID_KEY = 'client-id' +SEND_TELEMETRY_KEY = 'send-telemetry' ZONE_KEY = 'zone' KJOB_BATCH_IMAGE = 'batch-image' KJOB_BATCH_WORKING_DIRECTORY = 'batch-working-directory' @@ -47,6 +48,7 @@ CLUSTER_NAME_KEY, PROJECT_KEY, CLIENT_ID_KEY, + SEND_TELEMETRY_KEY, ZONE_KEY, GKE_ENDPOINT_KEY, DEPENDENCIES_KEY, diff --git a/src/xpk/core/telemetry.py b/src/xpk/core/telemetry.py index 216d2e82c..d49dd86ec 100644 --- a/src/xpk/core/telemetry.py +++ b/src/xpk/core/telemetry.py @@ -27,9 +27,17 @@ from enum import Enum from typing import Any from dataclasses import dataclass -from .config import xpk_config, CLIENT_ID_KEY, __version__ as xpk_version +from .config import xpk_config, CLIENT_ID_KEY, SEND_TELEMETRY_KEY, __version__ as xpk_version from ..utils.execution_context import is_dry_run from ..utils.user_agent import get_user_agent +from ..utils.feature_flags import FeatureFlags + + +def should_send_telemetry(): + return ( + FeatureFlags.TELEMETRY_ENABLED + and xpk_config.get(SEND_TELEMETRY_KEY) != "false" + ) def send_clearcut_payload(data: str, wait_to_complete: bool = False) -> None: diff --git a/src/xpk/core/telemetry_test.py b/src/xpk/core/telemetry_test.py index 488b03471..f22ae48b5 100644 --- a/src/xpk/core/telemetry_test.py +++ b/src/xpk/core/telemetry_test.py @@ -16,9 +16,10 @@ import pytest import json -from .config import xpk_config, CLIENT_ID_KEY -from .telemetry import MetricsCollector, MetricsEventMetadataKey +from .config import xpk_config, CLIENT_ID_KEY, SEND_TELEMETRY_KEY +from .telemetry import MetricsCollector, MetricsEventMetadataKey, should_send_telemetry from ..utils.execution_context import set_dry_run +from ..utils.feature_flags import FeatureFlags from pytest_mock import MockerFixture @@ -35,6 +36,23 @@ def setup_mocks(mocker: MockerFixture): xpk_config.set(CLIENT_ID_KEY, None) +@pytest.mark.parametrize( + argnames='feature_flag,config_value,expected', + argvalues=[ + (True, 'true', True), + (False, 'true', False), + (True, None, True), + (True, 'false', False), + ], +) +def test_should_send_telemetry_returns_correct_value( + feature_flag: bool, config_value: str, expected: bool +): + xpk_config.set(SEND_TELEMETRY_KEY, config_value) + FeatureFlags.TELEMETRY_ENABLED = feature_flag + assert should_send_telemetry() is expected + + def test_metrics_collector_generates_client_id_if_not_present(): xpk_config.set(CLIENT_ID_KEY, None) MetricsCollector.log_start(command='test') diff --git a/src/xpk/main.py b/src/xpk/main.py index 069f188de..0ad77675c 100644 --- a/src/xpk/main.py +++ b/src/xpk/main.py @@ -37,8 +37,7 @@ from .parser.core import set_parser from .core.updates import print_xpk_hello -from .core.telemetry import MetricsCollector, send_clearcut_payload -from .utils.feature_flags import FeatureFlags +from .core.telemetry import MetricsCollector, send_clearcut_payload, should_send_telemetry from .utils.console import xpk_print, exit_code_to_int from .utils.execution_context import set_context ################### Compatibility Check ################### @@ -89,7 +88,7 @@ def main() -> None: MetricsCollector.log_complete(-1) raise finally: - if FeatureFlags.TELEMETRY_ENABLED: + if should_send_telemetry(): send_clearcut_payload(MetricsCollector.flush())