From aa56abddfe9cb7558fbafd0b96364569e86cab46 Mon Sep 17 00:00:00 2001
From: Arthur Mesquita Pickcius <apickcius@esss.co>
Date: Mon, 27 Mar 2023 11:00:34 -0300
Subject: [PATCH] Add EnumKeySerializer tests

RFDAP-5724
---
 src/serialchemy/_tests/sample_model.py        |  4 ++--
 src/serialchemy/_tests/test_serialization.py  | 22 +++++++++++++++++--
 .../test_enum_key_field_dump.yml              | 20 +++++++++++++++++
 .../test_enum_key_field_load.yml              | 12 ++++++++++
 src/serialchemy/enum_field.py                 |  2 +-
 5 files changed, 55 insertions(+), 5 deletions(-)
 create mode 100644 src/serialchemy/_tests/test_serialization/test_enum_key_field_dump.yml
 create mode 100644 src/serialchemy/_tests/test_serialization/test_enum_key_field_load.yml

diff --git a/src/serialchemy/_tests/sample_model.py b/src/serialchemy/_tests/sample_model.py
index b414092..61c9cf9 100644
--- a/src/serialchemy/_tests/sample_model.py
+++ b/src/serialchemy/_tests/sample_model.py
@@ -8,7 +8,7 @@
 from sqlalchemy.sql import sqltypes
 from sqlalchemy_utils import ChoiceType
 
-from serialchemy.enum_serializer import EnumKeySerializer
+from serialchemy.enum_field import EnumKeyField
 from serialchemy.field import Field
 from serialchemy.model_serializer import ModelSerializer
 from serialchemy.nested_fields import NestedModelField, NestedModelListField
@@ -165,4 +165,4 @@ class EmployeeSerializer(ModelSerializer):
     company_name = Field(dump_only=True)
     address = NestedModelField(Address)
     contacts = NestedModelListField(Contact)
-    marital_status = Field(serializer=EnumKeySerializer(MaritalStatus))
+    marital_status = EnumKeyField(MaritalStatus)
diff --git a/src/serialchemy/_tests/test_serialization.py b/src/serialchemy/_tests/test_serialization.py
index 0a0c8d4..ad540ed 100644
--- a/src/serialchemy/_tests/test_serialization.py
+++ b/src/serialchemy/_tests/test_serialization.py
@@ -1,8 +1,8 @@
 import pytest
 
 from serialchemy._tests.sample_model import (
-    Address, Company, ContractType, Department, Employee, Engineer, Manager, MaritalStatus,
-    SpecialistEngineer)
+    Address, Company, ContractType, Department, Employee, EmployeeSerializer, Engineer, Manager,
+    MaritalStatus, SpecialistEngineer)
 from serialchemy.field import Field
 from serialchemy.func import dump
 from serialchemy.model_serializer import ModelSerializer
@@ -79,6 +79,12 @@ def test_model_dump(db_session, data_regression):
     serialized = serializer.dump(emp)
     data_regression.Check(serialized)
 
+def test_enum_key_field_dump(db_session, data_regression):
+    emp = db_session.query(Employee).get(1)
+    serializer = EmployeeSerializer(Employee)
+    serialized = serializer.dump(emp)
+    data_regression.Check(serialized)
+
 
 def test_model_load(data_regression):
     serializer = ModelSerializer(Employee)
@@ -92,6 +98,18 @@ def test_model_load(data_regression):
     model = serializer.load(employee_dict)
     data_regression.Check(dump(model))
 
+def test_enum_key_field_load(data_regression):
+    serializer = EmployeeSerializer(Employee)
+    employee_dict = {
+        "firstname": "Sarah",
+        "lastname": "Kerrigan",
+        "email": "sarahk@blitz.com",
+        "admission": "2152-01-02T00:00:00",
+        "marital_status": "MARRIED"
+    }
+    model = serializer.load(employee_dict)
+    data_regression.Check(dump(model))
+
 
 def test_one2one_pk_field(db_session, data_regression):
     serializer = EmployeeSerializerPrimaryKeyFields(Employee)
diff --git a/src/serialchemy/_tests/test_serialization/test_enum_key_field_dump.yml b/src/serialchemy/_tests/test_serialization/test_enum_key_field_dump.yml
new file mode 100644
index 0000000..3b73c21
--- /dev/null
+++ b/src/serialchemy/_tests/test_serialization/test_enum_key_field_dump.yml
@@ -0,0 +1,20 @@
+address:
+  city: Tarsonis
+  id: 1
+  number: '943'
+  state: null
+  street: 5 Av
+  zip: null
+address_id: 1
+admission: '2000-01-01'
+company_id: 5
+company_name: Terrans
+contacts: []
+contract_type: null
+created_at: '2000-01-02T00:00:00'
+email: null
+firstname: Jim
+id: 1
+lastname: Raynor
+marital_status: MARRIED
+role: Manager
diff --git a/src/serialchemy/_tests/test_serialization/test_enum_key_field_load.yml b/src/serialchemy/_tests/test_serialization/test_enum_key_field_load.yml
new file mode 100644
index 0000000..24d073d
--- /dev/null
+++ b/src/serialchemy/_tests/test_serialization/test_enum_key_field_load.yml
@@ -0,0 +1,12 @@
+address_id: null
+admission: '2152-01-02'
+company_id: null
+contract_type: null
+created_at: null
+email: sarahk@blitz.com
+firstname: Sarah
+id: null
+lastname: Kerrigan
+marital_status: Married
+password: null
+role: Employee
diff --git a/src/serialchemy/enum_field.py b/src/serialchemy/enum_field.py
index 4c31771..e2968c8 100644
--- a/src/serialchemy/enum_field.py
+++ b/src/serialchemy/enum_field.py
@@ -4,4 +4,4 @@
 
 class EnumKeyField(Field):
     def __init__(self, enum_class, dump_only=False, load_only=False, creation_only=False):
-        super().__init__(dump_only, load_only, creation_only, EnumKeySerializer(enum_class))
\ No newline at end of file
+        super().__init__(dump_only=dump_only, load_only=load_only, creation_only=creation_only, serializer=EnumKeySerializer(enum_class))
\ No newline at end of file