test_metric.py

import os
import unittest

from unittest.mock import patch, call

from botocore.exceptions import ClientError as BotocoreClientError
from datadog.api.exceptions import ClientError
from datetime import datetime, timedelta

from datadog_lambda.metric import lambda_metric, flush_stats
from datadog_lambda.api import decrypt_kms_api_key, KMS_ENCRYPTION_CONTEXT_KEY
from datadog_lambda.thread_stats_writer import ThreadStatsWriter
from datadog_lambda.tags import dd_lambda_layer_tag


class TestLambdaMetric(unittest.TestCase):
    def setUp(self):
        patcher = patch("datadog_lambda.metric.lambda_stats")
        self.mock_metric_lambda_stats = patcher.start()
        self.addCleanup(patcher.stop)

    def test_lambda_metric_tagged_with_dd_lambda_layer(self):
        lambda_metric("test", 1)
        lambda_metric("test", 1, 123, [])
        lambda_metric("test", 1, tags=["tag1:test"])
        self.mock_metric_lambda_stats.distribution.assert_has_calls(
            [
                call("test", 1, timestamp=None, tags=[dd_lambda_layer_tag]),
                call("test", 1, timestamp=123, tags=[dd_lambda_layer_tag]),
                call(
                    "test", 1, timestamp=None, tags=["tag1:test", dd_lambda_layer_tag]
                ),
            ]
        )

    # let's fake that the extension is present, this should override DD_FLUSH_TO_LOG
    @patch("datadog_lambda.metric.should_use_extension", True)
    def test_lambda_metric_flush_to_log_with_extension(self):
        os.environ["DD_FLUSH_TO_LOG"] = "True"
        lambda_metric("test", 1)
        self.mock_metric_lambda_stats.distribution.assert_has_calls(
            [call("test", 1, timestamp=None, tags=[dd_lambda_layer_tag])]
        )
        del os.environ["DD_FLUSH_TO_LOG"]

    @patch("datadog_lambda.metric.should_use_extension", True)
    def test_lambda_metric_timestamp_with_extension(self):
        patcher = patch("datadog_lambda.metric.extension_thread_stats")
        self.mock_metric_extension_thread_stats = patcher.start()
        self.addCleanup(patcher.stop)

        delta = timedelta(minutes=1)
        timestamp = int((datetime.now() - delta).timestamp())

        lambda_metric("test_timestamp", 1, timestamp)
        self.mock_metric_lambda_stats.distribution.assert_not_called()
        self.mock_metric_extension_thread_stats.distribution.assert_called_with(
            "test_timestamp", 1, timestamp=timestamp, tags=[dd_lambda_layer_tag]
        )

    @patch("datadog_lambda.metric.should_use_extension", True)
    def test_lambda_metric_datetime_with_extension(self):
        patcher = patch("datadog_lambda.metric.extension_thread_stats")
        self.mock_metric_extension_thread_stats = patcher.start()
        self.addCleanup(patcher.stop)

        delta = timedelta(hours=5)
        timestamp = datetime.now() - delta

        lambda_metric("test_timestamp", 1, timestamp)
        self.mock_metric_lambda_stats.distribution.assert_not_called()
        self.mock_metric_extension_thread_stats.distribution.assert_not_called()

    @patch("datadog_lambda.metric.should_use_extension", True)
    def test_lambda_metric_invalid_timestamp_with_extension(self):
        patcher = patch("datadog_lambda.metric.extension_thread_stats")
        self.mock_metric_extension_thread_stats = patcher.start()
        self.addCleanup(patcher.stop)

        delta = timedelta(hours=5)
        timestamp = int((datetime.now() - delta).timestamp())

        lambda_metric("test_timestamp", 1, timestamp)
        self.mock_metric_lambda_stats.distribution.assert_not_called()
        self.mock_metric_extension_thread_stats.distribution.assert_not_called()

    def test_lambda_metric_flush_to_log(self):
        os.environ["DD_FLUSH_TO_LOG"] = "True"

        lambda_metric("test", 1)
        self.mock_metric_lambda_stats.distribution.assert_not_called()

        del os.environ["DD_FLUSH_TO_LOG"]

    @patch("datadog_lambda.metric.logger.warning")
    def test_lambda_metric_invalid_metric_name_none(self, mock_logger_warning):
        lambda_metric(None, 1)
        self.mock_metric_lambda_stats.distribution.assert_not_called()
        mock_logger_warning.assert_called_once_with(
            "Ignoring metric submission. Invalid metric name: %s", None
        )

    @patch("datadog_lambda.metric.logger.warning")
    def test_lambda_metric_invalid_metric_name_not_string(self, mock_logger_warning):
        lambda_metric(123, 1)
        self.mock_metric_lambda_stats.distribution.assert_not_called()
        mock_logger_warning.assert_called_once_with(
            "Ignoring metric submission. Invalid metric name: %s", 123
        )

    @patch("datadog_lambda.metric.logger.warning")
    def test_lambda_metric_non_numeric_value(self, mock_logger_warning):
        lambda_metric("test.non_numeric", "oops")
        self.mock_metric_lambda_stats.distribution.assert_not_called()
        mock_logger_warning.assert_called_once_with(
            "Ignoring metric submission for metric '%s' because the value is not numeric: %r",
            "test.non_numeric",
            "oops",
        )


class TestFlushThreadStats(unittest.TestCase):
    def setUp(self):
        patcher = patch(
            "datadog.threadstats.reporters.HttpReporter.flush_distributions"
        )
        self.mock_threadstats_flush_distributions = patcher.start()
        self.addCleanup(patcher.stop)

        patcher = patch("datadog_lambda.metric.extension_thread_stats")
        self.mock_extension_thread_stats = patcher.start()
        self.addCleanup(patcher.stop)

    def test_retry_on_remote_disconnected(self):
        # Raise the RemoteDisconnected error
        lambda_stats = ThreadStatsWriter(True)

        self.mock_threadstats_flush_distributions.side_effect = ClientError(
            "POST",
            "https://api.datadoghq.com/api/v1/distribution_points",
            "RemoteDisconnected('Remote end closed connection without response')",
        )
        lambda_stats.flush()
        self.assertEqual(self.mock_threadstats_flush_distributions.call_count, 2)

    def test_flush_stats_with_tags(self):
        lambda_stats = ThreadStatsWriter(True)
        original_constant_tags = lambda_stats.thread_stats.constant_tags.copy()
        tags = ["tag1:value1", "tag2:value2"]

        # Add a metric to be flushed
        lambda_stats.distribution("test.metric", 1, tags=["metric:tag"])

        with patch.object(
            lambda_stats.thread_stats.reporter, "flush_distributions"
        ) as mock_flush_distributions:
            lambda_stats.flush(tags)
            mock_flush_distributions.assert_called_once()
            # Verify that after flush, constant_tags is reset to original
            self.assertEqual(
                lambda_stats.thread_stats.constant_tags, original_constant_tags
            )

    def test_flush_temp_constant_tags(self):
        lambda_stats = ThreadStatsWriter(flush_in_thread=True)
        lambda_stats.thread_stats.constant_tags = ["initial:tag"]
        original_constant_tags = lambda_stats.thread_stats.constant_tags.copy()

        lambda_stats.distribution("test.metric", 1, tags=["metric:tag"])
        flush_tags = ["flush:tag1", "flush:tag2"]

        with patch.object(
            lambda_stats.thread_stats.reporter, "flush_distributions"
        ) as mock_flush_distributions:
            lambda_stats.flush(tags=flush_tags)
            mock_flush_distributions.assert_called_once()
            flushed_dists = mock_flush_distributions.call_args[0][0]

            # Expected tags: original constant_tags + flush_tags + metric tags
            expected_tags = original_constant_tags + flush_tags + ["metric:tag"]

            # Verify the tags on the metric
            self.assertEqual(len(flushed_dists), 1)
            metric = flushed_dists[0]
            self.assertEqual(sorted(metric["tags"]), sorted(expected_tags))

            # Verify that constant_tags is reset after flush
            self.assertEqual(
                lambda_stats.thread_stats.constant_tags, original_constant_tags
            )

        # Repeat to ensure tags do not accumulate over multiple flushes
        new_flush_tags = ["flush:tag3"]
        lambda_stats.distribution("test.metric2", 2, tags=["metric2:tag"])

        with patch.object(
            lambda_stats.thread_stats.reporter, "flush_distributions"
        ) as mock_flush_distributions:
            lambda_stats.flush(tags=new_flush_tags)
            mock_flush_distributions.assert_called_once()
            flushed_dists = mock_flush_distributions.call_args[0][0]
            # Expected tags for the new metric
            expected_tags = original_constant_tags + new_flush_tags + ["metric2:tag"]

            self.assertEqual(len(flushed_dists), 1)
            metric = flushed_dists[0]
            self.assertEqual(sorted(metric["tags"]), sorted(expected_tags))
            self.assertEqual(
                lambda_stats.thread_stats.constant_tags, original_constant_tags
            )

    def test_flush_stats_without_context(self):
        flush_stats(lambda_context=None)
        self.mock_extension_thread_stats.flush.assert_called_with(None)


MOCK_FUNCTION_NAME = "myFunction"

# An API key encrypted with KMS and encoded as a base64 string
MOCK_ENCRYPTED_API_KEY_BASE64 = "MjIyMjIyMjIyMjIyMjIyMg=="

# The encrypted API key after it has been decoded from base64
MOCK_ENCRYPTED_API_KEY = "2222222222222222"

# The true value of the API key after decryption by KMS
EXPECTED_DECRYPTED_API_KEY = "1111111111111111"


class TestDecryptKMSApiKey(unittest.TestCase):
    def test_key_encrypted_with_encryption_context(self):
        os.environ["AWS_LAMBDA_FUNCTION_NAME"] = MOCK_FUNCTION_NAME

        class MockKMSClient:
            def decrypt(self, CiphertextBlob=None, EncryptionContext={}):
                if (
                    EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY)
                    != MOCK_FUNCTION_NAME
                ):
                    raise BotocoreClientError({}, "Decrypt")
                if CiphertextBlob == MOCK_ENCRYPTED_API_KEY.encode("utf-8"):
                    return {
                        "Plaintext": EXPECTED_DECRYPTED_API_KEY.encode("utf-8"),
                    }

        mock_kms_client = MockKMSClient()
        decrypted_key = decrypt_kms_api_key(
            mock_kms_client, MOCK_ENCRYPTED_API_KEY_BASE64
        )
        self.assertEqual(decrypted_key, EXPECTED_DECRYPTED_API_KEY)

        del os.environ["AWS_LAMBDA_FUNCTION_NAME"]

    def test_key_encrypted_without_encryption_context(self):
        class MockKMSClient:
            def decrypt(self, CiphertextBlob=None, EncryptionContext={}):
                if EncryptionContext.get(KMS_ENCRYPTION_CONTEXT_KEY) != None:
                    raise BotocoreClientError({}, "Decrypt")
                if CiphertextBlob == MOCK_ENCRYPTED_API_KEY.encode("utf-8"):
                    return {
                        "Plaintext": EXPECTED_DECRYPTED_API_KEY.encode("utf-8"),
                    }

        mock_kms_client = MockKMSClient()
        decrypted_key = decrypt_kms_api_key(
            mock_kms_client, MOCK_ENCRYPTED_API_KEY_BASE64
        )
        self.assertEqual(decrypted_key, EXPECTED_DECRYPTED_API_KEY)