From cbcd45fee1a8117888cce9f245637dea902ea8e1 Mon Sep 17 00:00:00 2001
From: Khayo CDM <84685469+khayomacedo@users.noreply.github.com>
Date: Thu, 26 Sep 2024 14:45:17 -0300
Subject: [PATCH] Add Helpers build_messages_attributes (#243)

---
 faster_sam/dependencies/events.py |  8 +++++-
 faster_sam/helpers.py             | 17 +++++++++++
 tests/test_helpers.py             | 47 +++++++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 1 deletion(-)
 create mode 100644 faster_sam/helpers.py
 create mode 100644 tests/test_helpers.py

diff --git a/faster_sam/dependencies/events.py b/faster_sam/dependencies/events.py
index 45a7d29..372d13f 100644
--- a/faster_sam/dependencies/events.py
+++ b/faster_sam/dependencies/events.py
@@ -4,6 +4,7 @@
 from typing import Any, Callable, Dict, Type
 from uuid import uuid4
 import uuid
+from faster_sam import helpers
 
 from fastapi import Request
 from pydantic import BaseModel
@@ -44,6 +45,11 @@ def dep(message: schema) -> Dict[str, Any]:
         assert isinstance(message, IntoSQSInfo)
 
         info = message.into()
+
+        message_attributes = {}
+        if info.message_attributes:
+            message_attributes = helpers.build_message_attributes(info.message_attributes)
+
         event = {
             "Records": [
                 {
@@ -56,7 +62,7 @@ def dep(message: schema) -> Dict[str, Any]:
                         "SenderId": str(uuid.uuid4()),
                         "ApproximateFirstReceiveTimestamp": info.sent_timestamp,
                     },
-                    "messageAttributes": info.message_attributes,
+                    "messageAttributes": message_attributes,
                     "md5OfBody": hashlib.md5(info.body.encode()).hexdigest(),
                     "eventSource": "aws:sqs",
                     "eventSourceARN": info.source_arn,
diff --git a/faster_sam/helpers.py b/faster_sam/helpers.py
new file mode 100644
index 0000000..f9a3dd3
--- /dev/null
+++ b/faster_sam/helpers.py
@@ -0,0 +1,17 @@
+from typing import Dict
+
+
+def build_message_attributes(attributes: Dict[str, str]) -> Dict[str, Dict[str, str]]:
+    attrs = {}
+
+    for key, value in attributes.items():
+        if isinstance(value, str):
+            attrs[key] = {"stringValue": value, "dataType": "String"}
+        elif isinstance(value, (int, float)):
+            attrs[key] = {"stringValue": str(value), "dataType": "Number"}
+        elif isinstance(value, bytes):
+            attrs[key] = {"BinaryValue": str(value), "dataType": "Binary"}
+        else:
+            raise TypeError(f"{value} of type {type(value).__name__} is not supported")
+
+    return attrs
diff --git a/tests/test_helpers.py b/tests/test_helpers.py
new file mode 100644
index 0000000..ea00891
--- /dev/null
+++ b/tests/test_helpers.py
@@ -0,0 +1,47 @@
+import unittest
+
+from faster_sam import helpers
+
+
+class TestHelpers(unittest.TestCase):
+    def teste_empty_attributes(self):
+        result = helpers.build_message_attributes({})
+        self.assertEqual(result, {})
+
+    def test_string_attributes(self):
+        attributes = {"key1": "value1", "key2": "value2"}
+        expected_result = {
+            "key1": {"stringValue": "value1", "dataType": "String"},
+            "key2": {"stringValue": "value2", "dataType": "String"},
+        }
+        result = helpers.build_message_attributes(attributes)
+        self.assertEqual(result, expected_result)
+
+    def test_number_attributes(self):
+        attributes = {"key1": 123, "key2": 12.3}
+        expected_result = {
+            "key1": {"stringValue": "123", "dataType": "Number"},
+            "key2": {"stringValue": "12.3", "dataType": "Number"},
+        }
+        result = helpers.build_message_attributes(attributes)
+        self.assertEqual(result, expected_result)
+
+    def test_bytes_attributes(self):
+        attributes = {"key1": b"test_bytes"}
+        expected_result = {"key1": {"BinaryValue": str(b"test_bytes"), "dataType": "Binary"}}
+        result = helpers.build_message_attributes(attributes)
+        self.assertEqual(result, expected_result)
+
+    def test_multi_attributes(self):
+        attributes = {"key1": 123, "key2": "value1"}
+        expected_result = {
+            "key1": {"stringValue": "123", "dataType": "Number"},
+            "key2": {"stringValue": "value1", "dataType": "String"},
+        }
+        result = helpers.build_message_attributes(attributes)
+        self.assertEqual(result, expected_result)
+
+    def test_unsupported_type(self):
+        attributes = {"key1": "value1", "key2": [1, 2, 3]}
+        with self.assertRaises(TypeError):
+            helpers.build_message_attributes(attributes)