From 4c47eae70bd8b30063e284b01a297b35bd70d618 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Fri, 24 Jun 2022 12:39:47 -0400
Subject: [PATCH] Common: Encapsulate AgentConfigurationSchema

---
 monkey/common/configuration/__init__.py             |  1 -
 .../configuration/test_agent_configuration.py       | 13 ++++---------
 2 files changed, 4 insertions(+), 10 deletions(-)

diff --git a/monkey/common/configuration/__init__.py b/monkey/common/configuration/__init__.py
index 107e5a491..1505f0e68 100644
--- a/monkey/common/configuration/__init__.py
+++ b/monkey/common/configuration/__init__.py
@@ -1,6 +1,5 @@
 from .agent_configuration import (
     AgentConfiguration,
-    AgentConfigurationSchema,
 )
 from .agent_sub_configurations import (
     CustomPBAConfiguration,
diff --git a/monkey/tests/unit_tests/common/configuration/test_agent_configuration.py b/monkey/tests/unit_tests/common/configuration/test_agent_configuration.py
index 23830d619..0de34e4a2 100644
--- a/monkey/tests/unit_tests/common/configuration/test_agent_configuration.py
+++ b/monkey/tests/unit_tests/common/configuration/test_agent_configuration.py
@@ -25,11 +25,8 @@ from tests.common.example_agent_configuration import (
     WINDOWS_FILENAME,
 )
 
-from common.configuration import (
-    DEFAULT_AGENT_CONFIGURATION_JSON,
-    AgentConfiguration,
-    AgentConfigurationSchema,
-)
+from common.configuration import DEFAULT_AGENT_CONFIGURATION_JSON, AgentConfiguration
+from common.configuration.agent_configuration import AgentConfigurationSchema
 from common.configuration.agent_sub_configuration_schemas import (
     CustomPBAConfigurationSchema,
     ExploitationConfigurationSchema,
@@ -159,10 +156,8 @@ def test_propagation_configuration():
 
 
 def test_agent_configuration():
-    schema = AgentConfigurationSchema()
-
     config = AgentConfiguration.from_dict(AGENT_CONFIGURATION)
-    config_dict = schema.dump(config)
+    config_json = AgentConfiguration.to_json(config)
 
     assert isinstance(config, AgentConfiguration)
     assert config.keep_tunnel_open_time == 30
@@ -171,7 +166,7 @@ def test_agent_configuration():
     assert isinstance(config.credential_collectors[0], PluginConfiguration)
     assert isinstance(config.payloads[0], PluginConfiguration)
     assert isinstance(config.propagation, PropagationConfiguration)
-    assert config_dict == AGENT_CONFIGURATION
+    assert json.loads(config_json) == AGENT_CONFIGURATION
 
 
 def test_default_agent_configuration():