nearly got tests passing - need to get correct item from list

3 files changed, 31 insertions, 2 deletions

diff --git a/datamaps/core/temporal.py b/datamaps/core/temporal.py
index d368f6a..23d4b23 100644
--- a/datamaps/core/temporal.py
+++ b/datamaps/core/temporal.py
@@ -1,5 +1,6 @@
-import datetime
 import calendar
+import datetime
+import itertools
 
 
 class Month:
@@ -146,6 +147,22 @@ class Quarter:
 
         self.start_date = self._start_date(self.quarter, self.year)
         self.end_date = self._end_date(self.quarter, self.year)
+        self.months = []
+        self._populate_months()
+
+    def _populate_months(self):
+        months_ints = []
+        start_int = Quarter._start_months[self.quarter][0]
+        strt = itertools.count(start_int)
+        for x in range(3):
+            months_ints.append(next(strt))
+        for m in months_ints:
+            if self.quarter == 4:
+                year = self.year + 1
+            else:
+                year = self.year
+            self.months.append(Month(m, year))
+
 
     def __str__(self):
         return f"Q{self.quarter} {str(self.year)[2:]}/{str(self.year + 1)[2:]}"
diff --git a/datamaps/plugins/dft/master.py b/datamaps/plugins/dft/master.py
index c3e6c3b..1eab1b7 100644
--- a/datamaps/plugins/dft/master.py
+++ b/datamaps/plugins/dft/master.py
@@ -162,7 +162,8 @@ class Master:
         """
         Returns the ``Month`` object associated with the ``Master``.
         """
-        pass
+        breakpoint()
+        return [m.month for m in self.quarter.months]
 
 
     @property
diff --git a/datamaps/tests/test_api.py b/datamaps/tests/test_api.py
index 21ab4d1..347ea4a 100644
--- a/datamaps/tests/test_api.py
+++ b/datamaps/tests/test_api.py
@@ -1,4 +1,5 @@
 import datetime
+import pytest
 
 from ..api import project_data_from_master, project_data_from_master_month
 from ..core.temporal import Quarter, Month
@@ -25,7 +26,17 @@ def test_get_project_data_using_month(master):
 
 def test_quarter_objects_have_months():
     q = Quarter(1, 2021)
+    q2 = Quarter(4, 2021)
     assert q.months[0].start_date == datetime.date(2021, 4, 1)
+    assert q.months[1].start_date == datetime.date(2021, 5, 1)
+    assert q.months[2].start_date == datetime.date(2021, 6, 1)
+    with pytest.raises(IndexError):
+        q.months[4].start_date == datetime.date(2021, 6, 1)
+    assert q2.months[0].start_date == datetime.date(2022, 1, 1)
+    assert q2.months[1].start_date == datetime.date(2022, 2, 1)
+    assert q2.months[2].start_date == datetime.date(2022, 3, 1)
+    with pytest.raises(IndexError):
+        q2.months[4].start_date == datetime.date(2022, 6, 1)
 
 
 def test_month():