Avoid conditional check on foreign key

1 files changed, 12 insertions, 2 deletions

diff --git a/yaksh/models.py b/yaksh/models.py
index 072c1b5..065d9cd 100644
--- a/yaksh/models.py
+++ b/yaksh/models.py
@@ -598,7 +598,12 @@ class LearningUnit(models.Model):
     check_prerequisite = models.BooleanField(default=True)
 
     def get_lesson_or_quiz(self):
-        return self.lesson if self.lesson else self.quiz
+        unit = None
+        if self.type == 'lesson':
+            unit = self.lesson
+        else:
+            unit = self.quiz
+        return unit
 
     def toggle_check_prerequisite(self):
         if self.check_prerequisite:
@@ -649,7 +654,12 @@ class LearningUnit(models.Model):
         return new_unit
 
     def __str__(self):
-        return self.lesson.name if self.lesson else self.quiz.description
+        name = None
+        if self.type == 'lesson':
+            name = self.lesson.name
+        else:
+            name = self.quiz.description
+        return name
 
 
 ###############################################################################