Bug: feat(themes): Setup DB for themes



There should a separate table for themes for storing the `theme_id`s for a particular lineage. This should be the schema:

| org_id | merchant_id | profile_id | theme_id |
| ------ | ----------- | ---------- | -------- |
| o1     | m1          | p1         | t1       |
| o1     | m1          | p2         | t2       |


This will be used for changing theme automatically when user switches any entities.
diff --git a/crates/common_utils/src/types.rs b/crates/common_utils/src/types.rs
index ef7a4b847c4..f7cdfd3617b 100644
--- a/crates/common_utils/src/types.rs
+++ b/crates/common_utils/src/types.rs
@@ -3,6 +3,8 @@ pub mod keymanager;
 
 /// Enum for Authentication Level
 pub mod authentication;
+/// Enum for Theme Lineage
+pub mod theme;
 
 use std::{
     borrow::Cow,
diff --git a/crates/common_utils/src/types/theme.rs b/crates/common_utils/src/types/theme.rs
new file mode 100644
index 00000000000..a2e6fe4b19c
--- /dev/null
+++ b/crates/common_utils/src/types/theme.rs
@@ -0,0 +1,39 @@
+use crate::id_type;
+
+/// Enum for having all the required lineage for every level.
+/// Currently being used for theme related APIs and queries.
+#[derive(Debug)]
+pub enum ThemeLineage {
+    /// Tenant lineage variant
+    Tenant {
+        /// tenant_id: String
+        tenant_id: String,
+    },
+    /// Org lineage variant
+    Organization {
+        /// tenant_id: String
+        tenant_id: String,
+        /// org_id: OrganizationId
+        org_id: id_type::OrganizationId,
+    },
+    /// Merchant lineage variant
+    Merchant {
+        /// tenant_id: String
+        tenant_id: String,
+        /// org_id: OrganizationId
+        org_id: id_type::OrganizationId,
+        /// merchant_id: MerchantId
+        merchant_id: id_type::MerchantId,
+    },
+    /// Profile lineage variant
+    Profile {
+        /// tenant_id: String
+        tenant_id: String,
+        /// org_id: OrganizationId
+        org_id: id_type::OrganizationId,
+        /// merchant_id: MerchantId
+        merchant_id: id_type::MerchantId,
+        /// profile_id: ProfileId
+        profile_id: id_type::ProfileId,
+    },
+}
diff --git a/crates/diesel_models/src/query/user.rs b/crates/diesel_models/src/query/user.rs
index 2bd403a847b..1f6e16702fb 100644
--- a/crates/diesel_models/src/query/user.rs
+++ b/crates/diesel_models/src/query/user.rs
@@ -1,6 +1,8 @@
 use common_utils::pii;
 use diesel::{associations::HasTable, ExpressionMethods};
+
 pub mod sample_data;
+pub mod theme;
 
 use crate::{
     query::generics, schema::users::dsl as users_dsl, user::*, PgPooledConn, StorageResult,
diff --git a/crates/diesel_models/src/query/user/theme.rs b/crates/diesel_models/src/query/user/theme.rs
new file mode 100644
index 00000000000..c021edca325
--- /dev/null
+++ b/crates/diesel_models/src/query/user/theme.rs
@@ -0,0 +1,95 @@
+use common_utils::types::theme::ThemeLineage;
+use diesel::{
+    associations::HasTable,
+    pg::Pg,
+    sql_types::{Bool, Nullable},
+    BoolExpressionMethods, ExpressionMethods, NullableExpressionMethods,
+};
+
+use crate::{
+    query::generics,
+    schema::themes::dsl,
+    user::theme::{Theme, ThemeNew},
+    PgPooledConn, StorageResult,
+};
+
+impl ThemeNew {
+    pub async fn insert(self, conn: &PgPooledConn) -> StorageResult<Theme> {
+        generics::generic_insert(conn, self).await
+    }
+}
+
+impl Theme {
+    fn lineage_filter(
+        lineage: ThemeLineage,
+    ) -> Box<
+        dyn diesel::BoxableExpression<<Self as HasTable>::Table, Pg, SqlType = Nullable<Bool>>
+            + 'static,
+    > {
+        match lineage {
+            ThemeLineage::Tenant { tenant_id } => Box::new(
+                dsl::tenant_id
+                    .eq(tenant_id)
+                    .and(dsl::org_id.is_null())
+                    .and(dsl::merchant_id.is_null())
+                    .and(dsl::profile_id.is_null())
+                    .nullable(),
+            ),
+            ThemeLineage::Organization { tenant_id, org_id } => Box::new(
+                dsl::tenant_id
+                    .eq(tenant_id)
+                    .and(dsl::org_id.eq(org_id))
+                    .and(dsl::merchant_id.is_null())
+                    .and(dsl::profile_id.is_null()),
+            ),
+            ThemeLineage::Merchant {
+                tenant_id,
+                org_id,
+                merchant_id,
+            } => Box::new(
+                dsl::tenant_id
+                    .eq(tenant_id)
+                    .and(dsl::org_id.eq(org_id))
+                    .and(dsl::merchant_id.eq(merchant_id))
+                    .and(dsl::profile_id.is_null()),
+            ),
+            ThemeLineage::Profile {
+                tenant_id,
+                org_id,
+                merchant_id,
+                profile_id,
+            } => Box::new(
+                dsl::tenant_id
+                    .eq(tenant_id)
+                    .and(dsl::org_id.eq(org_id))
+                    .and(dsl::merchant_id.eq(merchant_id))
+                    .and(dsl::profile_id.eq(profile_id)),
+            ),
+        }
+    }
+
+    pub async fn find_by_lineage(
+        conn: &PgPooledConn,
+        lineage: ThemeLineage,
+    ) -> StorageResult<Self> {
+        generics::generic_find_one::<<Self as HasTable>::Table, _, _>(
+            conn,
+            Self::lineage_filter(lineage),
+        )
+        .await
+    }
+
+    pub async fn delete_by_theme_id_and_lineage(
+        conn: &PgPooledConn,
+        theme_id: String,
+        lineage: ThemeLineage,
+    ) -> StorageResult<Self> {
+        generics::generic_delete_one_with_result::<<Self as HasTable>::Table, _, _>(
+            conn,
+            dsl::theme_id
+                .eq(theme_id)
+                .and(Self::lineage_filter(lineage)),
+        )
+        .await
+    }
+}
diff --git a/crates/diesel_models/src/schema.rs b/crates/diesel_models/src/schema.rs
index e2ab676b2d3..17de01b1862 100644
--- a/crates/diesel_models/src/schema.rs
+++ b/crates/diesel_models/src/schema.rs
@@ -1262,6 +1262,26 @@ diesel::table! {
     }
 }
 
+diesel::table! {
+    use diesel::sql_types::*;
+    use crate::enums::diesel_exports::*;
+
+    themes (theme_id) {
+        #[max_length = 64]
+        theme_id -> Varchar,
+        #[max_length = 64]
+        tenant_id -> Varchar,
+        #[max_length = 64]
+        org_id -> Nullable<Varchar>,
+        #[max_length = 64]
+        merchant_id -> Nullable<Varchar>,
+        #[max_length = 64]
+        profile_id -> Nullable<Varchar>,
+        created_at -> Timestamp,
+        last_modified_at -> Timestamp,
+    }
+}
+
 diesel::table! {
     use diesel::sql_types::*;
     use crate::enums::diesel_exports::*;
@@ -1406,6 +1426,7 @@ diesel::allow_tables_to_appear_in_same_query!(
     reverse_lookup,
     roles,
     routing_algorithm,
+    themes,
     unified_translations,
     user_authentication_methods,
     user_key_store,
diff --git a/crates/diesel_models/src/schema_v2.rs b/crates/diesel_models/src/schema_v2.rs
index 3bd31f5cbf3..7b4ccee2c4b 100644
--- a/crates/diesel_models/src/schema_v2.rs
+++ b/crates/diesel_models/src/schema_v2.rs
@@ -1208,6 +1208,26 @@ diesel::table! {
     }
 }
 
+diesel::table! {
+    use diesel::sql_types::*;
+    use crate::enums::diesel_exports::*;
+
+    themes (theme_id) {
+        #[max_length = 64]
+        theme_id -> Varchar,
+        #[max_length = 64]
+        tenant_id -> Varchar,
+        #[max_length = 64]
+        org_id -> Nullable<Varchar>,
+        #[max_length = 64]
+        merchant_id -> Nullable<Varchar>,
+        #[max_length = 64]
+        profile_id -> Nullable<Varchar>,
+        created_at -> Timestamp,
+        last_modified_at -> Timestamp,
+    }
+}
+
 diesel::table! {
     use diesel::sql_types::*;
     use crate::enums::diesel_exports::*;
@@ -1353,6 +1373,7 @@ diesel::allow_tables_to_appear_in_same_query!(
     reverse_lookup,
     roles,
     routing_algorithm,
+    themes,
     unified_translations,
     user_authentication_methods,
     user_key_store,
diff --git a/crates/diesel_models/src/user.rs b/crates/diesel_models/src/user.rs
index 9f7b77dc5d6..cf584c09b12 100644
--- a/crates/diesel_models/src/user.rs
+++ b/crates/diesel_models/src/user.rs
@@ -6,8 +6,9 @@ use time::PrimitiveDateTime;
 use crate::{diesel_impl::OptionalDieselArray, enums::TotpStatus, schema::users};
 
 pub mod dashboard_metadata;
-
 pub mod sample_data;
+pub mod theme;
+
 #[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
 #[diesel(table_name = users, primary_key(user_id), check_for_backend(diesel::pg::Pg))]
 pub struct User {
diff --git a/crates/diesel_models/src/user/theme.rs b/crates/diesel_models/src/user/theme.rs
new file mode 100644
index 00000000000..0824ae71919
--- /dev/null
+++ b/crates/diesel_models/src/user/theme.rs
@@ -0,0 +1,29 @@
+use common_utils::id_type;
+use diesel::{Identifiable, Insertable, Queryable, Selectable};
+use time::PrimitiveDateTime;
+
+use crate::schema::themes;
+
+#[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
+#[diesel(table_name = themes, primary_key(theme_id), check_for_backend(diesel::pg::Pg))]
+pub struct Theme {
+    pub theme_id: String,
+    pub tenant_id: String,
+    pub org_id: Option<id_type::OrganizationId>,
+    pub merchant_id: Option<id_type::MerchantId>,
+    pub profile_id: Option<id_type::ProfileId>,
+    pub created_at: PrimitiveDateTime,
+    pub last_modified_at: PrimitiveDateTime,
+}
+
+#[derive(Clone, Debug, Insertable, router_derive::DebugAsDisplay)]
+#[diesel(table_name = themes)]
+pub struct ThemeNew {
+    pub theme_id: String,
+    pub tenant_id: String,
+    pub org_id: Option<id_type::OrganizationId>,
+    pub merchant_id: Option<id_type::MerchantId>,
+    pub profile_id: Option<id_type::ProfileId>,
+    pub created_at: PrimitiveDateTime,
+    pub last_modified_at: PrimitiveDateTime,
+}
diff --git a/crates/router/src/db.rs b/crates/router/src/db.rs
index 81319ae4c0b..377253e9735 100644
--- a/crates/router/src/db.rs
+++ b/crates/router/src/db.rs
@@ -145,6 +145,7 @@ pub trait GlobalStorageInterface:
     + dyn_clone::DynClone
     + user::UserInterface
     + user_key_store::UserKeyStoreInterface
+    + user::theme::ThemeInterface
     + 'static
 {
 }
diff --git a/crates/router/src/db/kafka_store.rs b/crates/router/src/db/kafka_store.rs
index 4b99a9c8f3a..d7d283819e5 100644
--- a/crates/router/src/db/kafka_store.rs
+++ b/crates/router/src/db/kafka_store.rs
@@ -1,7 +1,11 @@
 use std::sync::Arc;
 
 use common_enums::enums::MerchantStorageScheme;
-use common_utils::{errors::CustomResult, id_type, pii, types::keymanager::KeyManagerState};
+use common_utils::{
+    errors::CustomResult,
+    id_type, pii,
+    types::{keymanager::KeyManagerState, theme::ThemeLineage},
+};
 use diesel_models::{
     enums,
     enums::ProcessTrackerStatus,
@@ -34,7 +38,7 @@ use time::PrimitiveDateTime;
 use super::{
     dashboard_metadata::DashboardMetadataInterface,
     role::RoleInterface,
-    user::{sample_data::BatchSampleDataInterface, UserInterface},
+    user::{sample_data::BatchSampleDataInterface, theme::ThemeInterface, UserInterface},
     user_authentication_method::UserAuthenticationMethodInterface,
     user_key_store::UserKeyStoreInterface,
     user_role::{ListUserRolesByOrgIdPayload, ListUserRolesByUserIdPayload, UserRoleInterface},
@@ -3683,3 +3687,30 @@ impl UserAuthenticationMethodInterface for KafkaStore {
             .await
     }
 }
+
+#[async_trait::async_trait]
+impl ThemeInterface for KafkaStore {
+    async fn insert_theme(
+        &self,
+        theme: storage::theme::ThemeNew,
+    ) -> CustomResult<storage::theme::Theme, errors::StorageError> {
+        self.diesel_store.insert_theme(theme).await
+    }
+
+    async fn find_theme_by_lineage(
+        &self,
+        lineage: ThemeLineage,
+    ) -> CustomResult<storage::theme::Theme, errors::StorageError> {
+        self.diesel_store.find_theme_by_lineage(lineage).await
+    }
+
+    async fn delete_theme_by_lineage_and_theme_id(
+        &self,
+        theme_id: String,
+        lineage: ThemeLineage,
+    ) -> CustomResult<storage::theme::Theme, errors::StorageError> {
+        self.diesel_store
+            .delete_theme_by_lineage_and_theme_id(theme_id, lineage)
+            .await
+    }
+}
diff --git a/crates/router/src/db/user.rs b/crates/router/src/db/user.rs
index 3cf68551b52..14bed15fa45 100644
--- a/crates/router/src/db/user.rs
+++ b/crates/router/src/db/user.rs
@@ -11,6 +11,7 @@ use crate::{
     services::Store,
 };
 pub mod sample_data;
+pub mod theme;
 
 #[async_trait::async_trait]
 pub trait UserInterface {
diff --git a/crates/router/src/db/user/theme.rs b/crates/router/src/db/user/theme.rs
new file mode 100644
index 00000000000..d71b82cdea4
--- /dev/null
+++ b/crates/router/src/db/user/theme.rs
@@ -0,0 +1,203 @@
+use common_utils::types::theme::ThemeLineage;
+use diesel_models::user::theme as storage;
+use error_stack::report;
+
+use super::MockDb;
+use crate::{
+    connection,
+    core::errors::{self, CustomResult},
+    services::Store,
+};
+
+#[async_trait::async_trait]
+pub trait ThemeInterface {
+    async fn insert_theme(
+        &self,
+        theme: storage::ThemeNew,
+    ) -> CustomResult<storage::Theme, errors::StorageError>;
+
+    async fn find_theme_by_lineage(
+        &self,
+        lineage: ThemeLineage,
+    ) -> CustomResult<storage::Theme, errors::StorageError>;
+
+    async fn delete_theme_by_lineage_and_theme_id(
+        &self,
+        theme_id: String,
+        lineage: ThemeLineage,
+    ) -> CustomResult<storage::Theme, errors::StorageError>;
+}
+
+#[async_trait::async_trait]
+impl ThemeInterface for Store {
+    async fn insert_theme(
+        &self,
+        theme: storage::ThemeNew,
+    ) -> CustomResult<storage::Theme, errors::StorageError> {
+        let conn = connection::pg_connection_write(self).await?;
+        theme
+            .insert(&conn)
+            .await
+            .map_err(|error| report!(errors::StorageError::from(error)))
+    }
+
+    async fn find_theme_by_lineage(
+        &self,
+        lineage: ThemeLineage,
+    ) -> CustomResult<storage::Theme, errors::StorageError> {
+        let conn = connection::pg_connection_read(self).await?;
+        storage::Theme::find_by_lineage(&conn, lineage)
+            .await
+            .map_err(|error| report!(errors::StorageError::from(error)))
+    }
+
+    async fn delete_theme_by_lineage_and_theme_id(
+        &self,
+        theme_id: String,
+        lineage: ThemeLineage,
+    ) -> CustomResult<storage::Theme, errors::StorageError> {
+        let conn = connection::pg_connection_write(self).await?;
+        storage::Theme::delete_by_theme_id_and_lineage(&conn, theme_id, lineage)
+            .await
+            .map_err(|error| report!(errors::StorageError::from(error)))
+    }
+}
+
+fn check_theme_with_lineage(theme: &storage::Theme, lineage: &ThemeLineage) -> bool {
+    match lineage {
+        ThemeLineage::Tenant { tenant_id } => {
+            &theme.tenant_id == tenant_id
+                && theme.org_id.is_none()
+                && theme.merchant_id.is_none()
+                && theme.profile_id.is_none()
+        }
+        ThemeLineage::Organization { tenant_id, org_id } => {
+            &theme.tenant_id == tenant_id
+                && theme
+                    .org_id
+                    .as_ref()
+                    .is_some_and(|org_id_inner| org_id_inner == org_id)
+                && theme.merchant_id.is_none()
+                && theme.profile_id.is_none()
+        }
+        ThemeLineage::Merchant {
+            tenant_id,
+            org_id,
+            merchant_id,
+        } => {
+            &theme.tenant_id == tenant_id
+                && theme
+                    .org_id
+                    .as_ref()
+                    .is_some_and(|org_id_inner| org_id_inner == org_id)
+                && theme
+                    .merchant_id
+                    .as_ref()
+                    .is_some_and(|merchant_id_inner| merchant_id_inner == merchant_id)
+                && theme.profile_id.is_none()
+        }
+        ThemeLineage::Profile {
+            tenant_id,
+            org_id,
+            merchant_id,
+            profile_id,
+        } => {
+            &theme.tenant_id == tenant_id
+                && theme
+                    .org_id
+                    .as_ref()
+                    .is_some_and(|org_id_inner| org_id_inner == org_id)
+                && theme
+                    .merchant_id
+                    .as_ref()
+                    .is_some_and(|merchant_id_inner| merchant_id_inner == merchant_id)
+                && theme
+                    .profile_id
+                    .as_ref()
+                    .is_some_and(|profile_id_inner| profile_id_inner == profile_id)
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl ThemeInterface for MockDb {
+    async fn insert_theme(
+        &self,
+        new_theme: storage::ThemeNew,
+    ) -> CustomResult<storage::Theme, errors::StorageError> {
+        let mut themes = self.themes.lock().await;
+        for theme in themes.iter() {
+            if new_theme.theme_id == theme.theme_id {
+                return Err(errors::StorageError::DuplicateValue {
+                    entity: "theme_id",
+                    key: None,
+                }
+                .into());
+            }
+
+            if new_theme.tenant_id == theme.tenant_id
+                && new_theme.org_id == theme.org_id
+                && new_theme.merchant_id == theme.merchant_id
+                && new_theme.profile_id == theme.profile_id
+            {
+                return Err(errors::StorageError::DuplicateValue {
+                    entity: "lineage",
+                    key: None,
+                }
+                .into());
+            }
+        }
+
+        let theme = storage::Theme {
+            theme_id: new_theme.theme_id,
+            tenant_id: new_theme.tenant_id,
+            org_id: new_theme.org_id,
+            merchant_id: new_theme.merchant_id,
+            profile_id: new_theme.profile_id,
+            created_at: new_theme.created_at,
+            last_modified_at: new_theme.last_modified_at,
+        };
+        themes.push(theme.clone());
+
+        Ok(theme)
+    }
+
+    async fn find_theme_by_lineage(
+        &self,
+        lineage: ThemeLineage,
+    ) -> CustomResult<storage::Theme, errors::StorageError> {
+        let themes = self.themes.lock().await;
+        themes
+            .iter()
+            .find(|theme| check_theme_with_lineage(theme, &lineage))
+            .cloned()
+            .ok_or(
+                errors::StorageError::ValueNotFound(format!(
+                    "Theme with lineage {:?} not found",
+                    lineage
+                ))
+                .into(),
+            )
+    }
+
+    async fn delete_theme_by_lineage_and_theme_id(
+        &self,
+        theme_id: String,
+        lineage: ThemeLineage,
+    ) -> CustomResult<storage::Theme, errors::StorageError> {
+        let mut themes = self.themes.lock().await;
+        let index = themes
+            .iter()
+            .position(|theme| {
+                theme.theme_id == theme_id && check_theme_with_lineage(theme, &lineage)
+            })
+            .ok_or(errors::StorageError::ValueNotFound(format!(
+                "Theme with id {} and lineage {:?} not found",
+                theme_id, lineage
+            )))?;
+
+        let theme = themes.remove(index);
+
+        Ok(theme)
+    }
+}
diff --git a/crates/storage_impl/src/mock_db.rs b/crates/storage_impl/src/mock_db.rs
index b3358d898b2..efcce567727 100644
--- a/crates/storage_impl/src/mock_db.rs
+++ b/crates/storage_impl/src/mock_db.rs
@@ -60,6 +60,7 @@ pub struct MockDb {
     pub user_key_store: Arc<Mutex<Vec<store::user_key_store::UserKeyStore>>>,
     pub user_authentication_methods:
         Arc<Mutex<Vec<store::user_authentication_method::UserAuthenticationMethod>>>,
+    pub themes: Arc<Mutex<Vec<store::user::theme::Theme>>>,
 }
 
 impl MockDb {
@@ -105,6 +106,7 @@ impl MockDb {
             roles: Default::default(),
             user_key_store: Default::default(),
             user_authentication_methods: Default::default(),
+            themes: Default::default(),
         })
     }
 }
diff --git a/migrations/2024-11-06-121933_setup-themes-table/down.sql b/migrations/2024-11-06-121933_setup-themes-table/down.sql
new file mode 100644
index 00000000000..4b590c34705
--- /dev/null
+++ b/migrations/2024-11-06-121933_setup-themes-table/down.sql
@@ -0,0 +1,3 @@
+-- This file should undo anything in `up.sql`
+DROP INDEX IF EXISTS themes_index;
+DROP TABLE IF EXISTS themes;
diff --git a/migrations/2024-11-06-121933_setup-themes-table/up.sql b/migrations/2024-11-06-121933_setup-themes-table/up.sql
new file mode 100644
index 00000000000..3d84fcd8140
--- /dev/null
+++ b/migrations/2024-11-06-121933_setup-themes-table/up.sql
@@ -0,0 +1,17 @@
+-- Your SQL goes here
+CREATE TABLE IF NOT EXISTS themes (
+    theme_id VARCHAR(64) PRIMARY KEY,
+    tenant_id VARCHAR(64) NOT NULL,
+    org_id VARCHAR(64),
+    merchant_id VARCHAR(64),
+    profile_id VARCHAR(64),
+    created_at TIMESTAMP NOT NULL,
+    last_modified_at TIMESTAMP NOT NULL
+);
+
+CREATE UNIQUE INDEX IF NOT EXISTS themes_index ON themes (
+    tenant_id,
+    COALESCE(org_id, '0'),
+    COALESCE(merchant_id, '0'),
+    COALESCE(profile_id, '0')
+);

## Type of Change
<!-- Put an `x` in the boxes that apply -->

- [ ] Bugfix
- [x] New feature
- [ ] Enhancement
- [ ] Refactoring
- [ ] Dependency updates
- [ ] Documentation
- [ ] CI/CD

## Description
<!-- Describe your changes in detail -->
This PR creates a new table named `themes` in the DB. This will be used to store the `theme_id` for any lineage of `tenant_id`, `org_id`, `merchant_id` and `profile_id`.

### Additional Changes

- [ ] This PR modifies the API contract
- [x] This PR modifies the database schema
- [ ] This PR modifies application configuration/environment variables

<!--
Provide links to the files with corresponding changes.

Following are the paths where you can find config files:
1. `config`
2. `crates/router/src/configs`
3. `loadtest/config`
-->


## Motivation and Context
<!--
Why is this change required? What problem does it solve?
If it fixes an open issue, please link to the issue here.

If you don't have an issue, we'd recommend starting with one first so the PR
can focus on the implementation (unless it is an obvious bug or documentation fix
that will have little conversation).
-->
Closes #6532

## How did you test it?
<!--
Did you write an integration/unit/API test to verify the code changes?
Or did you test this change manually (provide relevant screenshots)?
-->
This is an internal change and there is no API to test this currently.

I've tested the read query by creating a temporary API which uses the query created in this PR.

## Checklist
<!-- Put an `x` in the boxes that apply -->

- [x] I formatted the code `cargo +nightly fmt --all`
- [x] I addressed lints thrown by `cargo clippy`
- [x] I reviewed the submitted code
- [ ] I added unit tests for my changes where possible
