Set highlight_letter_for_sentiment() back to returning a list

letters/tests/test_views.py

@@ -512,6 +512,19 @@ def test_get_highlighted_letter_sentiment(self, mock_highlight_letter_for_sentim
             self.assertTrue(item in content,
                             "get_highlighted_letter_sentiment() response content should contain sentiment value".format(item))
 
+        # If sentiment value is a list, all those values should end up in response
+        mock_highlight_letter_for_sentiment.return_value = [letter, letter]
+        sentiments = [('1', ['0.1234', '0.5678'])]
+
+        request = RequestFactory().get(reverse('letter_sentiment_view', kwargs={'letter_id': '1', 'sentiment_id': '1'}),
+                                       follow=True)
+        response = get_highlighted_letter_sentiment(request, letter, sentiments)
+        content = str(response.content)
+
+        for item in ['0.1234', '0.5678']:
+            self.assertTrue(item in content,
+                            "get_highlighted_letter_sentiment() response content should contain sentiment value".format(item))
+
 
 class HighlightLetterForSentimentTestCase(TestCase):
     """
@@ -524,6 +537,8 @@ def test_highlight_letter_for_sentiment(self, mock_body_as_text, mock_highlight_
         """
         highlight_letter_for_sentiment() should call highlight_for_sentiment()
         for each of a letter's fields and return a copy of the letter with fields highlighted
+
+        It returns a list because multiple sentiments get highlighted for positive/negative
         """
 
         mock_body_as_text.return_value = 'As this is the beginin of a new year I thought as I was a lone to night I ' \
@@ -557,7 +572,7 @@ def test_highlight_letter_for_sentiment(self, mock_body_as_text, mock_highlight_
         self.assertEqual(mock_highlight_for_sentiment.call_args_list[5][0], (letter.ps, 1),
                          'highlight_letter_for_sentiment() should call highlight_for_sentiment(ps, sentiment_id)')
 
-        self.assertEqual(letter.contents(), result.contents(),
+        self.assertEqual(letter.contents(), result[0].contents(),
                          'highlight_letter_for_sentiment() should return highlighted letter')
 
 

letters/views.py

@@ -285,7 +285,8 @@ def get_highlighted_letter_sentiment(request, letter, sentiments):
         else:
             sentiment_values.extend(value)
 
-        highlighted_letters.append(highlighted_letter)
+        # highlight_letter_for_sentiment always returns a list, so use extend
+        highlighted_letters.extend(highlighted_letter)
 
     results = zip(sentiment_values, highlighted_letters)
     return render(request, 'letter_sentiment.html',
@@ -294,26 +295,32 @@ def get_highlighted_letter_sentiment(request, letter, sentiments):
 
 
 def highlight_letter_for_sentiment(letter, sentiment_id):
-    # highlight_for_sentiment() returns a list, for the case of multiple different positive/negative
-    # sentiments. Just need the first result here
-    heading = highlight_for_sentiment(letter.heading, sentiment_id)[0] if letter.heading else ''
-    greeting = highlight_for_sentiment(letter.greeting, sentiment_id)[0] if letter.greeting else ''
-    body = highlight_for_sentiment(letter.body_as_text(), sentiment_id)[0] if letter.body else ''
-    closing = highlight_for_sentiment(letter.closing, sentiment_id)[0] if letter.closing else ''
-    sig = highlight_for_sentiment(letter.signature, sentiment_id)[0] if letter.signature else ''
-    ps = highlight_for_sentiment(letter.ps, sentiment_id)[0] if letter.ps else ''
-
-    # Make a copy of the letter so we can manipulate the content fields
-    highlighted_letter = deepcopy(letter)
-    highlighted_letter.pk = None
-    highlighted_letter.heading = mark_safe(heading)
-    highlighted_letter.greeting = mark_safe(greeting)
-    highlighted_letter.body = mark_safe(body)
-    highlighted_letter.closing = mark_safe(closing)
-    highlighted_letter.signature = mark_safe(sig)
-    highlighted_letter.ps = mark_safe(ps)
-
-    return highlighted_letter
+    # highlight_for_sentiment() returns a list,
+    # for the case of multiple different positive/negative sentiments
+    highlighted_letters = []
+
+    headings = highlight_for_sentiment(letter.heading, sentiment_id)
+    greetings = highlight_for_sentiment(letter.greeting, sentiment_id)
+    bodies = highlight_for_sentiment(letter.body_as_text(), sentiment_id)
+    closings = highlight_for_sentiment(letter.closing, sentiment_id)
+    sigs = highlight_for_sentiment(letter.signature, sentiment_id)
+    pss = highlight_for_sentiment(letter.ps, sentiment_id)
+
+    for idx, heading in enumerate(headings):
+        # Make a copy of the letter so we can manipulate the content fields
+        highlighted_letter = deepcopy(letter)
+        highlighted_letter.pk = None
+        highlighted_letter.heading = mark_safe(heading)
+        highlighted_letter.greeting = mark_safe(greetings[idx])
+        highlighted_letter.body = mark_safe(bodies[idx])
+        highlighted_letter.closing = mark_safe(closings[idx])
+        highlighted_letter.signature = mark_safe(sigs[idx])
+
+        highlighted_letter.ps = mark_safe(pss[idx])
+
+        highlighted_letters.append(highlighted_letter)
+
+    return highlighted_letters
 
 
 class TextSentimentView(TemplateView):