Made requested changes in models.py and test_models.py

- Remove 'category' key from legend_display_types dictionary.
- Add a condition to check if question type exists before creating category_question_map.
- Add test.
- Change lambda to list in defaultdict.

2 files changed, 16 insertions, 23 deletions

diff --git a/yaksh/models.py b/yaksh/models.py
index 0f20605..cce90e7 100644
--- a/yaksh/models.py
+++ b/yaksh/models.py
@@ -82,14 +82,14 @@ string_check_type = (
     )
 
 legend_display_types = {
-        "mcq": {"category": "objective", "label": "Objective Type"},
-        "mcc": {"category": "objective", "label": "Objective Type"},
-        "code": {"category": "programming", "label": "Programming"},
-        "upload": {"category": "upload", "label": "Upload"},
-        "integer": {"category": "objective", "label": "Objective Type"},
-        "string": {"category": "objective", "label": "Objective Type"},
-        "float": {"category": "objective", "label": "Objective Type"},
-        "arrange": {"category": "objective", "label": "Objective Type"},
+        "mcq": {"label": "Objective Type"},
+        "mcc": {"label": "Objective Type"},
+        "code": {"label": "Programming"},
+        "upload": {"label": "Upload"},
+        "integer": {"label": "Objective Type"},
+        "string": {"label": "Objective Type"},
+        "float": {"label": "Objective Type"},
+        "arrange": {"label": "Objective Type"},
     }
 
 attempts = [(i, i) for i in range(1, 6)]
@@ -2104,12 +2104,13 @@ class AnswerPaper(models.Model):
         return self.answers.filter(question=question).order_by('-id')
 
     def get_categorized_question_indices(self):
-
-        category_question_map = defaultdict(lambda: [])
+        category_question_map = defaultdict(list)
         for index, question in enumerate(self.get_all_ordered_questions(), 1):
-            category_question_map[
-                legend_display_types[question.type]["label"]
-            ].append(index)
+            question_category = legend_display_types.get(question.type)
+            if question_category:
+                category_question_map[
+                    question_category["label"]
+                ].append(index)
         return dict(category_question_map)
 
     def validate_answer(self, user_answer, question, json_data=None, uid=None,
diff --git a/yaksh/test_models.py b/yaksh/test_models.py
index 8b9aa37..2e136d5 100644
--- a/yaksh/test_models.py
+++ b/yaksh/test_models.py
@@ -1585,18 +1585,10 @@ class AnswerPaperTestCases(unittest.TestCase):
                 )
 
     def test_get_categorized_question_indices(self):
-        all_questions = self.answerpaper.get_all_ordered_questions()
+        question_indices = {'Programming': [1], 'Objective Type': [2, 3]}
         categorized_question_indices = \
             self.answerpaper.get_categorized_question_indices()
-        questions = all_questions[0:4]
-        category_question_map = defaultdict(lambda: [])
-        for index, question in enumerate(questions, 1):
-            category_question_map[
-                    legend_display_types[question.type]["label"]
-                    ].append(index)
-        category_question_map = dict(category_question_map)
-        self.assertDictEqual(
-                        category_question_map, categorized_question_indices)
+        self.assertDictEqual(question_indices, categorized_question_indices)
 
 
 ###############################################################################