From 7e3e2ba514285aff4e58153478f7412bed1387c1 Mon Sep 17 00:00:00 2001
From: GabrielBarberini <gabriel.barberini@sumup.com>
Date: Wed, 6 Mar 2024 19:07:56 -0300
Subject: [PATCH] refactors environment.set_date test

---
 tests/conftest.py              | 16 ++++++++++++----
 tests/unit/test_environment.py | 35 +++++++++++-----------------------
 2 files changed, 23 insertions(+), 28 deletions(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 9afcbbdd9..6c3db54c9 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1109,18 +1109,26 @@ def dimensionless_calisto(kg, m, dimensionless_cesaroni_m1670):
 
 @pytest.fixture
 def example_env():
-    """Create a simple object of the Environment class to be used in the tests.
-    This allows to avoid repeating the same code in all tests. The environment
-    set here is the simplest possible, with no parameters set.
+    """Simple object of the Environment class to be used in the tests.
 
     Returns
     -------
     rocketpy.Environment
-        The simplest object of the Environment class
     """
     return Environment()
 
 
+@pytest.fixture
+def example_date_naive():
+    """Naive tomorrow date
+
+    Returns
+    -------
+    datetime.datetime
+    """
+    return datetime.datetime.now()
+
+
 @pytest.fixture
 def example_env_robust():
     """Create an object of the Environment class to be used in the tests. This
diff --git a/tests/unit/test_environment.py b/tests/unit/test_environment.py
index 463238276..0b631a82b 100644
--- a/tests/unit/test_environment.py
+++ b/tests/unit/test_environment.py
@@ -1,4 +1,3 @@
-import datetime
 import os
 
 import numpy as np
@@ -8,41 +7,29 @@
 from rocketpy import Environment
 
 
-def test_env_set_date(example_env):
-    """Test that the date is set correctly in the environment object. This
-    basically takes a date and time and converts it to a datetime object, then
-    set the date to the example environment object. The test checks if the
-    datetime object is the same as the one in the example environment object.
+def test_date_naive_set_date(example_env, example_date_naive):
+    """Test that the time zone is set correctly to UTC by default in the environment object for a date_naive input
 
     Parameters
     ----------
     example_env : rocketpy.Environment
-        Example environment object to be tested.
+    example_date_naive: datetime.datetime
     """
-    tomorrow = datetime.date.today() + datetime.timedelta(days=1)
-    example_env.set_date((tomorrow.year, tomorrow.month, tomorrow.day, 12))
-    assert example_env.datetime_date == datetime.datetime(
-        tomorrow.year, tomorrow.month, tomorrow.day, 12, tzinfo=pytz.utc
-    )
+    example_env.set_date(example_date_naive)
+    assert example_env.datetime_date == pytz.utc.localize(example_date_naive)
 
 
-def test_env_set_date_time_zone(example_env):
-    """Test that the time zone is set correctly in the environment object
+def test_date_aware_set_date(example_env, example_date_naive):
+    """Test that the time zone is set accordingly in the environment object for a date_aware input
 
     Parameters
     ----------
     example_env : rocketpy.Environment
-        Example environment object to be tested.
+    example_date_naive: datetime.datetime
     """
-    tomorrow = datetime.date.today() + datetime.timedelta(days=1)
-    example_env.set_date(
-        (tomorrow.year, tomorrow.month, tomorrow.day, 12), timezone="America/New_York"
-    )
-    date_naive = datetime.datetime(tomorrow.year, tomorrow.month, tomorrow.day, 12)
-    timezone = pytz.timezone("America/New_York")
-    date_aware_local_date = timezone.localize(date_naive)
-    date_aware_utc = date_aware_local_date.astimezone(pytz.UTC)
-    assert example_env.datetime_date == date_aware_utc
+    example_env.set_date(example_date_naive, timezone="America/New_York")
+    example_date_aware = pytz.timezone("America/New_York").localize(example_date_naive)
+    assert example_env.datetime_date == example_date_aware
 
 
 def test_env_set_location(example_env):