From 9cc3639691399c7991aa73d91d1565faa8caf6c0 Mon Sep 17 00:00:00 2001
From: Dalton Barreto <daltonmatos@gmail.com>
Date: Wed, 25 Jan 2017 18:06:13 -0200
Subject: [PATCH] Adds MarathonApp.add_env() method

---
 marathon/models/app.py  |  5 ++++-
 tests/test_model_app.py | 26 ++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_model_app.py

diff --git a/marathon/models/app.py b/marathon/models/app.py
index 366e958..3053f55 100644
--- a/marathon/models/app.py
+++ b/marathon/models/app.py
@@ -119,7 +119,7 @@ def __init__(self, accepted_resource_roles=None, args=None, backoff_factor=None,
             for d in (deployments or [])
         ]
         self.disk = disk
-        self.env = env
+        self.env = env or dict()
         self.executor = executor
         self.gpus = gpus
         self.health_checks = health_checks or []
@@ -185,6 +185,9 @@ def __init__(self, accepted_resource_roles=None, args=None, backoff_factor=None,
         self.task_stats = task_stats if (isinstance(task_stats, MarathonTaskStats) or task_stats is None) \
             else MarathonTaskStats.from_json(task_stats)
 
+    def add_env(self, key, value):
+        self.env[key] = value
+
 
 class MarathonHealthCheck(MarathonObject):
 
diff --git a/tests/test_model_app.py b/tests/test_model_app.py
new file mode 100644
index 0000000..adddb4d
--- /dev/null
+++ b/tests/test_model_app.py
@@ -0,0 +1,26 @@
+# encoding: utf-8
+
+from marathon.models.app import MarathonApp
+import unittest
+
+
+class MarathonAppTest(unittest.TestCase):
+
+    def test_env_defaults_to_empty_dict(self):
+        """
+        é testé
+        """
+        app = MarathonApp()
+        self.assertEquals(app.env, {})
+
+    def test_add_env_empty_dict(self):
+        app = MarathonApp()
+        app.add_env("MY_ENV", "my-value")
+        self.assertDictEqual({"MY_ENV": "my-value"}, app.env)
+
+    def test_add_env_non_empty_dict(self):
+        env_data = {"OTHER_ENV": "other-value"}
+        app = MarathonApp(env=env_data)
+
+        app.add_env("MY_ENV", "my-value")
+        self.assertDictEqual({"MY_ENV": "my-value", "OTHER_ENV": "other-value"}, app.env)