Repository: Nike-Inc/koheesio

Issue:
[FEATURE] Bigger, better, SecretStr
We should create a custom SecretStr implementation that plays a bit nicer with str operations, but that's up for debate.

For example, this code would throw an error:

```python
from pydantic import SecretStr
body = SecretStr("foobarbazbladibla")
type = "bearer"
header = type + body
```
results in `TypeError: unsupported operand type(s) for +: 'SecretStr' and 'str'`

using f-string notation also doesn't work nicely:

```python
header = f"{type} {body}"
```
you don't get an error doing this, but the resulting type of header will be str and will be set to `'bearer **********'`. The resulting str also seems to be 'broken' in that you can't run type against it for example.

Long story short, we should make 'adding' things (prefixing of suffixes) something to a given SecretStr easy as well.

Discussion:
dannymeijer: I had to make some changes to properly accommodate Python 3.12 - please re-review

Patch:
diff --git a/pyproject.toml b/pyproject.toml
index 89641a3..ea59dfd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -410,7 +410,7 @@ Available scripts:
 - `coverage` or `cov` - run the test suite with coverage.
 """
 path = ".venv"
-python = "3.11"
+python = "3.12"
 template = "default"
 features = [
   "async",
diff --git a/src/koheesio/models/__init__.py b/src/koheesio/models/__init__.py
index 5482960..cad78f4 100644
--- a/src/koheesio/models/__init__.py
+++ b/src/koheesio/models/__init__.py
@@ -14,7 +14,10 @@ from __future__ import annotations
 from typing import Annotated, Any, Dict, List, Optional, Union
 from abc import ABC
 from functools import cached_property, partial
+import inspect
 from pathlib import Path
+import re
+import sys
 
 # to ensure that koheesio.models is a drop in replacement for pydantic
 from pydantic import BaseModel as PydanticBaseModel
@@ -27,8 +30,6 @@ from pydantic import (
     InstanceOf,
     PositiveInt,
     PrivateAttr,
-    SecretBytes,
-    SecretStr,
     SkipValidation,
     ValidationError,
     conint,
@@ -38,6 +39,8 @@ from pydantic import (
     field_validator,
     model_validator,
 )
+from pydantic import SecretBytes as PydanticSecretBytes
+from pydantic import SecretStr as PydanticSecretStr
 
 # noinspection PyProtectedMember
 from pydantic._internal._generics import PydanticGenericMetadata
@@ -764,6 +767,242 @@ In case an individual column is passed, the value will be coerced to a list.
 """
 
 
+class _SecretMixin:
+    """Mixin class that provides additional functionality to Pydantic's SecretStr and SecretBytes classes."""
+
+    def __add__(self, other: Union[str, bytes, '_SecretMixin']) -> '_SecretMixin':
+        """Support concatenation when the SecretMixin instance is on the left side of the + operator.
+        
+        Raises
+        ------
+        TypeError
+            If concatenation fails.
+        """
+        left = self.get_secret_value()
+        right = other.get_secret_value() if isinstance(other, _SecretMixin) else other
+        return self.__class__(left + right)  # type: ignore
+    
+    def __radd__(self, other: Union[str, bytes, '_SecretMixin']) -> '_SecretMixin':
+        """Support concatenation when the SecretMixin instance is on the right side of the + operator.
+
+        Raises
+        ------
+        TypeError
+            If concatenation fails.
+        """
+        right = self.get_secret_value()
+        left = other.get_secret_value() if isinstance(other, _SecretMixin) else other
+        return self.__class__(left + right)  # type: ignore
+    
+    def __mul__(self, n: int) -> '_SecretMixin':
+        """Support multiplication when the SecretMixin instance is on the left side of the * operator.
+
+        Raises
+        ------
+        TypeError
+            If multiplication fails.
+        """
+        if isinstance(n, int):
+            return self.__class__(self.get_secret_value() * n)  # type: ignore
+        return NotImplemented
+    
+    def __rmul__(self, n: int) -> '_SecretMixin':
+        """Support multiplication when the SecretMixin instance is on the right side of the * operator.
+
+        Raises
+        ------
+        TypeError
+            If multiplication fails.
+        """
+        return self.__mul__(n)
+
+
+class SecretStr(PydanticSecretStr, _SecretMixin):
+    """A string type that ensures the secrecy of its value, extending Pydantic's SecretStr.
+
+    This class provides additional functionality over Pydantic's SecretStr, including:
+    - Support for concatenation with other strings and SecretStr instances.
+    - Advanced f-string formatting support to ensure the secret value is only revealed in secure contexts.
+
+    For more information on Pydantic's SecretStr, see: https://docs.pydantic.dev/latest/usage/types/#secret-types
+
+    Examples
+    --------
+    ### Basic Usage
+    ```python
+    secret = SecretStr("my_secret")
+    ```
+
+    ### String representations of the secrets are masked
+    ```python
+    str(secret)
+    # '**********'
+    repr(secret)
+    # "SecretStr('**********')"
+    ```
+
+    ### Concatenations are supported with other strings and SecretStr instances
+    ```python
+    secret + "suffix"
+    # SecretStr('my_secretsuffix')
+    "prefix" + secret
+    # SecretStr('prefixmy_secret')
+    ```
+
+    ### f-string formatting is supported
+    If the f-string is called from within a SecretStr, the secret value is returned as we are in a secure context.
+    ```python
+    new_secret = f"{SecretStr(f'prefix{secret}suffix')}"
+    new_secret.get_secret_value()
+    ### 'prefixmy_secretsuffix'
+    ```
+    
+    Otherwise, we consider the context 'unsafe': the SecretStr instance is returned, and the secret value is masked.
+    ```python
+    f"{secret}"
+    '**********'
+    ```
+
+    Parameters
+    ----------
+    secret : str
+        The secret value to be stored.
+
+    Methods
+    -------
+    get_secret_value()
+        Returns the actual secret value.
+    """
+
+    def __format__(self, format_spec: str) -> str:
+        """Advanced f-string formatting support.
+        If the f-string is called from within a SecretStr, the secret value is returned as we are in a secure context.
+        Otherwise, we consider the context 'unsafe' and we let pydantic take care of the formatting.
+        """
+        # Inspect the call stack to determine if the string is being passed through SecretStr
+        stack = inspect.stack(context=1)
+        caller_context = stack[1].code_context[0]  # type: ignore
+
+        # f-string behavior is different from Python 3.12 onwards, so we need to handle it separately
+        if sys.version_info >= (3, 12):
+            caller_frame = stack[1].frame
+
+            # Check if we are in a multiline f-string
+            if not any(substring in caller_context for substring in ("f'", 'f"', '.format')):
+                # Multiline f-strings do not show the outer code context, we'll need to extract it from the source code
+                source_lines = inspect.getsourcelines(caller_frame)[0]
+                lineno = stack[1].positions.lineno  # the line of code that called this function
+
+                # Find the start of the multiline string f""" or f'''
+                starting_index = next(
+                    (
+                        lineno - i - 1
+                        for i, line in enumerate(reversed(source_lines[:lineno]))
+                        if any(marker in line for marker in ('f"""', "f'''"))
+                    ),
+                    None
+                )
+
+                # Extract the multiline string
+                multiline_string = "".join(source_lines[starting_index : lineno])
+
+                # Remove the code context from the multiline string
+                caller_context = multiline_string.replace(caller_context, "").strip()
+
+        # Remove comments from the caller context
+        caller_context = re.sub(r'#.*', '', caller_context).strip()
+
+        # Remove the entire string that the format method was called on: 
+        # 1. matches any string enclosed in single or double quotes that contains 'format('.
+        caller_context = re.sub(r'["\'].*?format\(.*?\)["\']', "", caller_context, flags=re.DOTALL).strip()
+        # 2. matches any string enclosed in single or double quotes, optionally prefixed with 'f' for f-strings.
+        caller_context = re.sub(r'f?["\'].*?["\']', "", caller_context, flags=re.DOTALL).strip()
+
+        # safe context: the secret value is returned
+        if 'SecretStr(' in caller_context:
+            return self.get_secret_value()
+
+        # unsafe context: we let pydantic handle the formatting
+        return super().__format__(format_spec)
+    
+
+class SecretBytes(PydanticSecretBytes, _SecretMixin):
+    """A bytes type that ensures the secrecy of its value, extending Pydantic's SecretBytes.
+
+    This class provides additional functionality over Pydantic's SecretBytes, including:
+    - Support for concatenation with other bytes and SecretBytes instances.
+    - Advanced f-string formatting support to ensure the secret value is only revealed in secure contexts.
+
+    For more information on Pydantic's SecretBytes, see: https://docs.pydantic.dev/latest/usage/types/#secret-types
+
+    Examples
+    --------
+    ### Basic Usage
+    ```python
+    secret = SecretBytes(b"my_secret")
+    ```
+
+    ### String representations of the secrets are masked
+    ```python
+    str(secret)
+    # '**********'
+    repr(secret)
+    # "SecretBytes('**********')"
+    ```
+
+    ### Concatenations are supported with other bytes and SecretBytes instances
+    ```python
+    secret + b"suffix"
+    # SecretBytes(b'my_secretsuffix')
+    b"prefix" + secret
+    # SecretBytes(b'prefixmy_secret')
+    ```
+
+    ### Multiplications are supported with integers
+    ```python
+    SecretBytes(b"foo") * 2
+    # SecretBytes(b'foofoo')
+    ```
+
+    Note
+    ----
+    The pydantic `SecretBytes` class is quite permissive with the types of data it accepts. Hence, you have to ensure 
+    that the data you pass to the `SecretBytes` '+' operator is of the right type to be concatenated - the `bytes` type
+    of the data is not guaranteed.
+
+    For example, python can 'add' two lists together like this:
+    
+    ```python
+    [1, 2] + [3, 4]  # [1, 2, 3, 4]
+    ```
+
+    If one of the list is wrapped in a SecretBytes instance, the '+' operator will work, just like the example above:
+    
+    ```python
+    list1 = SecretBytes([1,2,3])
+    list2 = [4,5,6]
+    list1 + list2  # SecretBytes([1, 2, 3, 4, 5, 6])
+    ```
+
+    If however you try to add any other type of data to the `list1` from the example above, you will get a `TypeError`.
+    You can use the `secret_data_type` method to check the type of the secret data stored in the `SecretBytes` instance.
+
+    Parameters
+    ----------
+    secret : bytes
+        The secret value to be stored.
+
+    Methods
+    -------
+    get_secret_value()
+        Returns the actual secret value.
+    """
+    
+    def secret_data_type(self) -> type:
+        """Return the type of the secret data."""
+        return type(self.get_secret_value())
+
+      
 def _list_of_strings_validation(strings_value: Union[str, list]) -> list:
     """
     Performs validation for ListOfStrings type. Will ensure that whether one string is provided or a list of strings,


Test patch:
diff --git a/tests/models/test_models.py b/tests/models/test_models.py
index 1a0374e..b386b98 100644
--- a/tests/models/test_models.py
+++ b/tests/models/test_models.py
@@ -1,17 +1,21 @@
 """Test suite for Koheesio's extended BaseModel class"""
 
-from typing import Optional
+from typing import Any, Optional
 import json
 from textwrap import dedent
 
 import pytest
 import yaml
 
+from pydantic import SecretBytes as PydanticSecretBytes
+from pydantic import SecretStr as PydanticSecretStr
+
 from koheesio.context import Context
-from koheesio.models import BaseModel, ExtraParamsMixin, ListOfStrings
+from koheesio.models import BaseModel, ExtraParamsMixin, SecretBytes, SecretStr, ListOfStrings
 
 
 class TestBaseModel:
+    """Test suite for BaseModel class"""
     class SimpleModel(BaseModel):
         a: int
         b: str = "default"
@@ -35,8 +39,8 @@ class TestBaseModel:
 
     def test_a_simple_model(self) -> None:
         """Test a simple model."""
-        foo = self.SimpleModel(a=1)
-        assert foo.model_dump() == {"a": 1, "b": "default", "description": "SimpleModel", "name": "SimpleModel"}
+        model = self.SimpleModel(a=1)
+        assert model.model_dump() == {"a": 1, "b": "default", "description": "SimpleModel", "name": "SimpleModel"}
 
     def test_context_management_no_exception(self) -> None:
         """Test that with-statement works without throwing exceptions"""
@@ -59,44 +63,55 @@ class TestBaseModel:
 
     @pytest.fixture(params=[{"foo": "bar"}, {"baz": 123}, {"foo": "bar", "baz": 123}])
     def context_data(self, request: pytest.FixtureRequest) -> dict:
+        """Fixture for context data"""
         return request.param
 
     def test_add(self) -> None:
+        """Test BaseModel's add method"""
+        # Arrange
         model1 = self.SimpleModel(a=1)
         model2 = self.SimpleModel(a=2)
+        # Act
         model = model1 + model2
+        # Assert
         assert isinstance(model, BaseModel)
         assert model.a == 2
         assert model.b == "default"
 
     def test_getitem(self) -> None:
+        """Test BaseModel's __getitem__ method"""
         model = self.SimpleModel(a=1)
         assert model["a"] == 1
 
     def test_setitem(self) -> None:
+        """Test BaseModel's __setitem__ method"""
         model = self.SimpleModel(a=1)
         model["a"] = 2
         assert model.a == 2
 
     def test_hasattr(self) -> None:
+        """Test BaseModel's hasattr method"""
         model = self.SimpleModel(a=1)
         assert model.hasattr("a")
         assert not model.hasattr("non_existent_key")
 
     def test_from_context(self, context_data: pytest.FixtureRequest) -> None:
+        """Test BaseModel's from_context method"""
         context = Context(context_data)
         model = self.FooModel.from_context(context)
         assert isinstance(model, BaseModel)
         for key, value in context_data.items():
             assert getattr(model, key) == value
 
-    def test_from_dict(self, context_data: pytest.FixtureRequest) -> None:
+    def test_from_dict(self, context_data: pytest.FixtureRequest) ->  None:
+        """Test BaseModel's from_dict method"""
         model = self.FooModel.from_dict(context_data)
         assert isinstance(model, BaseModel)
         for key, value in context_data.items():
             assert getattr(model, key) == value
 
     def test_from_json(self, context_data: pytest.FixtureRequest) -> None:
+        """Test BaseModel's from_json method"""
         json_data = json.dumps(context_data)
         model = self.FooModel.from_json(json_data)
         assert isinstance(model, BaseModel)
@@ -104,54 +119,76 @@ class TestBaseModel:
             assert getattr(model, key) == value
 
     def test_from_toml(self) -> None:
+        """Test BaseModel's from_toml method"""
+        # Arrange
         toml_data = dedent(
             """
             a = 1
             b = "default"
             """
         )
+        # Act
         model = self.SimpleModel.from_toml(toml_data)
+        # Assert
         assert isinstance(model, BaseModel)
         assert model.a == 1
         assert model.b == "default"
 
     def test_from_yaml(self, context_data: pytest.FixtureRequest) -> None:
+        """Test BaseModel's from_yaml method"""
+        # Arrange
         yaml_data = yaml.dump(context_data)
+        # Act
         model = self.FooModel.from_yaml(yaml_data)
+        # Assert
         assert isinstance(model, BaseModel)
         for key, value in context_data.items():
             assert getattr(model, key) == value
 
     def test_to_context(self) -> None:
+        """Test BaseModel's to_context method"""
+        # Arrange
         model = self.SimpleModel(a=1)
+        # Act
         context = model.to_context()
+        # Assert
         assert isinstance(context, Context)
         assert context.a == 1
         assert context.b == "default"
 
     def test_to_dict(self) -> None:
+        """Test BaseModel's to_dict method"""
+        # Arrange
         model = self.SimpleModel(a=1)
+        # Act
         dict_model = model.to_dict()
+        # Assert
         assert isinstance(dict_model, dict)
         assert dict_model["a"] == 1
         assert dict_model["b"] == "default"
 
     def test_to_json(self) -> None:
+        """Test BaseModel's to_json method"""
+        # Arrange
         model = self.SimpleModel(a=1)
+        # Act
         json_model = model.to_json()
+        # Assert
         assert isinstance(json_model, str)
         assert '"a": 1' in json_model
         assert '"b": "default"' in json_model
 
     def test_to_yaml(self) -> None:
+        """Test BaseModel's to_yaml method"""
+        # Arrange
         model = self.SimpleModel(a=1)
+        # Act
         yaml_model = model.to_yaml()
+        # Assert
         assert isinstance(yaml_model, str)
         assert "a: 1" in yaml_model
         assert "b: default" in yaml_model
 
-    import pytest
-
     class ModelWithDescription(BaseModel):
         a: int = 42
         description: str = "This is a\nmultiline description"
@@ -205,9 +242,9 @@ class TestBaseModel:
             ),
         ],
     )
-    def test_name_and_multiline_description(
-        self, model_class: type[BaseModel], instance_arg: dict, expected: dict
-    ) -> None:
+
+    def test_name_and_multiline_description(self, model_class: type[BaseModel], instance_arg: dict, expected: dict) -> None:
+        """Test that the name and description are correctly set."""
         instance = model_class(**instance_arg)
         assert instance.model_dump() == expected
 
@@ -226,9 +263,9 @@ class TestBaseModel:
             (ModelWithLongDescriptionAndNoSpaces, 120, "ThisIsAVeryLongDescription" * 4 + "ThisIsAVeryLo..."),
         ],
     )
-    def test_extremely_long_description(
-        self, model_class: type[BaseModel], expected_length: int, expected_description: str
-    ) -> None:
+
+    def test_extremely_long_description(self, model_class: type[BaseModel], expected_length: int, expected_description: str) -> None:
+        """Test that the description is truncated if it is too long."""
         model = model_class()
         assert len(model.description) == expected_length
         assert model.description == expected_description
@@ -236,7 +273,9 @@ class TestBaseModel:
 
 
 class TestExtraParamsMixin:
+    """Test suite for ExtraParamsMixin class"""
     def test_extra_params_mixin(self) -> None:
+        """Test ExtraParamsMixin class."""
         class SimpleModelWithExtraParams(BaseModel, ExtraParamsMixin):
             a: int
             b: str = "default"
@@ -253,6 +292,299 @@ class TestExtraParamsMixin:
         }
 
 
+class TestSecretStr:
+    """Test suite for SecretStr class"""
+    # reference values
+    secret_value = "foobarbazbladibla"
+    pydantic_secret = PydanticSecretStr(secret_value)
+    koheesio_secret = SecretStr(secret_value)
+    prefix = "prefix"
+    suffix = "suffix"
+
+    class StrMethodRaisesException:
+        """a class that does not implement __str__ raise a TypeError"""
+        def __str__(self):  # type: ignore
+            raise TypeError("Cannot convert to string")
+    
+    def test_secret_str_str(self) -> None:
+        """check that the integrity of the str method in SecretStr is preserved
+        by comparing Pydantic's SecretStr with Koheesio's SecretStr str method output"""
+        pydantic_secret_str = str(self.pydantic_secret)
+        actual = str(self.koheesio_secret)
+        expected = "**********"
+        assert pydantic_secret_str == actual == expected
+    
+    def test_secret_str_repr(self) -> None:
+        """check that the integrity of the repr method in SecretStr is preserved
+        by comparing Pydantic's SecretStr with Koheesio's SecretStr repr method output"""
+        pydantic_secret_str = repr(self.pydantic_secret)
+        actual = repr(self.koheesio_secret)
+        expected = "SecretStr('**********')"
+        assert pydantic_secret_str == actual == expected
+
+    @pytest.mark.parametrize("other", [
+        42,  # int
+        3.14,  # float
+        True,  # bool
+        None,  # None
+        [1, 2, 3],  # list
+        {"key": "value"},  # dict
+        (1, 2),  # tuple
+        {1, 2, 3},  # set
+        bytes("byte_string", "utf-8"),  # bytes
+        StrMethodRaisesException()  # custom class that raises TypeError in __str__
+    ])
+    def test_concatenate_unhappy(self, other: Any) -> None:
+        """check that concatenating a SecretStr with a non-stringable objects raises an exception"""
+        with pytest.raises(TypeError):
+            _ = self.koheesio_secret + other
+
+    def test_secret_str_with_f_string_secretstr(self) -> None:
+        """check that a str and SecretStr can be combined with one another using f-strings
+        Test through using f-string with a SecretStr. Here we expect that the secret gets properly processed.
+        """
+        # arrange
+        secret = self.koheesio_secret
+        # act
+        actual_secret = SecretStr(f"{self.prefix}{secret}{self.suffix}")
+        # assert
+        expected = PydanticSecretStr(f"{self.prefix}{self.secret_value}{self.suffix}")
+        assert actual_secret.get_secret_value() == expected.get_secret_value()
+
+    def test_secret_str_with_f_string_str(self) -> None:
+        """check that a str and SecretStr can be combined with one another using f-strings
+        Test through using f-string with a str. Here we expect the secret to remain hidden.
+        """
+        # arrange
+        secret = self.koheesio_secret
+        # act
+        actual_str = f"{self.prefix}{secret}{self.suffix}"
+        # assert
+        expected = f"{self.prefix}**********{self.suffix}"
+        assert actual_str == expected
+
+    def test_secret_str_add(self) -> None:
+        """check that a SecretStr and a str can be combined with one another using concatenation"""
+        # arrange
+        secret = self.koheesio_secret
+        # act
+        actual = secret + self.suffix
+        # assert
+        expected = PydanticSecretStr(self.secret_value + self.suffix)
+        assert actual.get_secret_value() == expected.get_secret_value()
+
+    def test_secret_str_radd(self) -> None:
+        """check that a str and SecretStr can be combined with one another using concatenation"""
+        # arrange
+        secret = self.koheesio_secret
+        # act
+        actual = self.prefix + secret
+        # assert
+        expected = PydanticSecretStr(self.prefix + self.secret_value)
+        assert actual.get_secret_value() == expected.get_secret_value()
+    
+    def test_add_two_secret_str(self) -> None:
+        """check that two SecretStr can be added together"""
+        # arrange
+        secret = self.koheesio_secret
+        # act
+        actual = secret + secret
+        # assert
+        expected = PydanticSecretStr(self.secret_value *2)
+        assert actual.get_secret_value() == expected.get_secret_value()
+
+    def test_secret_str_mul_and_rmul(self) -> None:
+        """check that a SecretBytes can be multiplied by an integer"""
+        # arrange
+        secret = self.koheesio_secret
+        # act
+        actual_mul = secret * 3
+        actual_rmul = 3 * secret
+        # assert
+        expected = PydanticSecretStr(self.secret_value * 3)
+        assert actual_mul.get_secret_value() == actual_rmul.get_secret_value() == expected.get_secret_value()
+    
+    @pytest.mark.parametrize(
+        "secret_value, format_spec, expected",
+        [
+            # Secure context with different format specifications
+            pytest.param("my_secret", "", "my_secret", id="secure context"),
+            pytest.param("my_secret", "s", "my_secret", id="secure context with format spec"),
+            pytest.param("my_secret", "r", "my_secret", id="secure context with !r"),
+            pytest.param("my_secret", "a", "my_secret", id="secure context with !a"),
+            # Empty secret value
+            pytest.param("", "", "", id="empty string"),
+            pytest.param("", "s", "", id="empty string with format spec"),
+            # Special characters in secret value
+            pytest.param("special_chars!@#", "", "special_chars!@#", id="special characters"),
+            pytest.param("special_chars!@#", "s", "special_chars!@#", id="special characters with format spec"),
+            # Multiline secret value
+            pytest.param("line1\nline2", "", "line1\nline2", id="multiline string"),
+            pytest.param("line1\nline2", "s", "line1\nline2", id="multiline string with format spec"),
+        ],
+    )
+    def test_secretstr_format_spec_secure_context(self, secret_value: str, format_spec: str, expected: str) -> None:
+        """Test that SecretStr can be formatted with a format spec"""
+        secret = SecretStr(secret_value)
+        actual = SecretStr(f"{secret:{format_spec}}")
+        assert actual.get_secret_value() == expected
+
+    def test_secretstr_format_edge_cases(self) -> None:
+        """Test that SecretStr remains secure in edge case situations
+        Note that all tests in here are not parameterized as the inline comments are crucial to the test cases.
+        """
+        input_str = SecretStr("my_super_secret")
+
+        # Test that SecretStr remains secure when an inline comment is added
+        non_secure_context = "This is a secret: " + input_str # This is a comment with SecretStr(
+        assert str(non_secure_context) == "**********"
+
+        # Test the same, but with using f-string
+        non_secure_context = f"This is a secret: {input_str}" # This is a comment with SecretStr(
+        assert str(non_secure_context) == "This is a secret: **********"
+
+        # Test string interpolation (non secure context)
+        interpolated_str = "{}".format(input_str)
+        assert interpolated_str == "**********"
+
+        # Test string interpolation (secure context)
+        interpolated_str = SecretStr("foo_{}".format(input_str))
+        assert str(interpolated_str) == "**********"
+        assert interpolated_str.get_secret_value() == "foo_my_super_secret"
+
+        # Test multiline interpolation (non secure context)
+        non_secure_context = """
+        foo.SecretStr({input_str})
+        """.format(input_str=input_str)
+        assert non_secure_context == "\n        foo.SecretStr(**********)\n        "
+
+        # Test multiline interpolation (secure context)
+        secure_context = SecretStr("""
+        foo.SecretStr({input_str})
+        """.format(input_str=input_str))
+        assert str(secure_context) == "**********"
+
+        # Ensure that we have no leakage of the secret value when using string interpolation
+        interpolated_str = "foo.SecretStr({})".format(input_str)
+        assert interpolated_str == "foo.SecretStr(**********)"
+
+        # Check for """ notation
+        non_secure_context = f"""foo.SecretStr({input_str})"""
+        assert non_secure_context == "foo.SecretStr(**********)"
+
+        # Check for ''' notation
+        non_secure_context = f'''foo.SecretStr({input_str})'''
+        assert non_secure_context == "foo.SecretStr(**********)"
+
+        # Check multiline f-string - non secure context
+        non_secure_context = f"""
+        foo.SecretStr({input_str})
+        """
+        assert non_secure_context == "\n        foo.SecretStr(**********)\n        "
+
+        # Check multiline f-string - secure context
+        secure_context = SecretStr(f"""
+        foo.{input_str}
+        """)
+        assert str(secure_context) == "**********"
+
+        # Check with nested SecretStr - secure context
+        secure_context = SecretStr(f"foo.{input_str}.{SecretStr('bar')}")
+        assert str(secure_context) == "**********"
+        assert secure_context.get_secret_value() == "foo.my_super_secret.bar"
+
+        # Check with nested SecretStr - non secure context
+        non_secure_context = f"foo.{input_str}.{SecretStr('bar')}"
+        assert non_secure_context == "foo.**********.**********"
+
+        # f-string with special characters - secure context
+        special_chars_f_string = SecretStr(f"!@#$%^&*({input_str})")
+        assert special_chars_f_string.get_secret_value() == "!@#$%^&*(my_super_secret)"
+
+        # f-string with special characters - non secure context
+        non_secure_context = f"!@#$%^&*({input_str})"
+        assert non_secure_context == "!@#$%^&*(**********)"
+
+        # f-string with escaped characters - secure context
+        escaped_chars_f_string = SecretStr(f"\\n{input_str}\\t")
+        assert escaped_chars_f_string.get_secret_value() == "\\nmy_super_secret\\t"
+
+        # f-string with escaped characters - non secure context
+        non_secure_context = f"\\n{input_str}\\t"
+        assert non_secure_context == "\\n**********\\t"
+
+class TestSecretBytes:
+    """Test suite for SecretBytes class"""
+    # reference values
+    secret_value = b"foobarbazbladibla"
+    pydantic_secret = PydanticSecretBytes(secret_value)
+    koheesio_secret = SecretBytes(secret_value)
+    prefix = b"prefix"
+    suffix = b"suffix"
+
+    def test_secret_bytes_str(self) -> None:
+        """check that the str method in SecretBytes is preserved"""
+        secret = self.koheesio_secret
+        actual = str(secret)
+        expected = "b'**********'"
+        assert actual == expected
+    
+    def test_secret_bytes_repr(self) -> None:
+        """check that the repr method in SecretBytes is preserve"""
+        secret = self.koheesio_secret
+        actual = repr(secret)
+        expected = "SecretBytes(b'**********')"
+        assert actual == expected
+
+    def test_secret_bytes_add(self) -> None:
+        """check that a SecretBytes and a bytes can be combined with one another using concatenation"""
+        # arrange
+        secret = self.koheesio_secret
+        # act
+        actual = secret + self.suffix
+        # assert
+        expected = PydanticSecretBytes(self.secret_value + self.suffix)
+        assert actual.get_secret_value() == expected.get_secret_value()
+
+    def test_secret_bytes_radd(self) -> None:
+        """check that a bytes and SecretBytes can be combined with one another using concatenation"""
+        # arrange
+        secret = self.koheesio_secret
+        # act
+        actual = self.prefix + secret
+        # assert
+        expected = PydanticSecretBytes(self.prefix + self.secret_value)
+        assert actual.get_secret_value() == expected.get_secret_value()
+
+    def test_add_two_secret_bytes(self) -> None:
+        """check that two SecretBytes can be added together"""
+        # arrange
+        secret = self.koheesio_secret
+        # act
+        actual = secret + secret
+        # assert
+        expected = PydanticSecretBytes(self.secret_value *2)
+        assert actual.get_secret_value() == expected.get_secret_value()
+
+    def test_secret_bytes_mul_and_rmul(self) -> None:
+        """check that a SecretBytes can be multiplied by an integer"""
+        # arrange
+        secret = self.koheesio_secret
+        # act
+        actual_mul = secret * 3
+        actual_rmul = 3 * secret
+        # assert
+        expected = PydanticSecretBytes(self.secret_value * 3)
+        assert actual_mul.get_secret_value() == actual_rmul.get_secret_value() == expected.get_secret_value()
+
+    def test_secret_data_type(self) -> None:
+        """check that the correct type is returned. Pydantic's SecretBytes maintains the data type passed to it"""
+        # arrange
+        secret = SecretBytes([1, 2, 3])
+        # act and assert
+        assert isinstance(secret.get_secret_value(), list)
+
+
 class TestAnnotatedTypes:
     class SomeModelWithListOfStrings(BaseModel):
         a: ListOfStrings
@@ -269,3 +601,4 @@ class TestAnnotatedTypes:
     def test_list_of_strings(self, list_of_strings, expected_list_of_strings) -> None:
         model_with = TestAnnotatedTypes.SomeModelWithListOfStrings(a=list_of_strings)
         assert model_with.a == expected_list_of_strings
+
diff --git a/tests/sso/test_okta.py b/tests/sso/test_okta.py
index d980ca1..10fc943 100644
--- a/tests/sso/test_okta.py
+++ b/tests/sso/test_okta.py
@@ -4,8 +4,7 @@ import logging
 import pytest
 from requests_mock.mocker import Mocker
 
-from pydantic import SecretStr
-
+from koheesio.models import SecretStr
 from koheesio.sso import okta as o