From c53704cf7cc8c2ce9d508d92a37eb480b998e32f Mon Sep 17 00:00:00 2001
From: Rob Hallam <0504004h@student.gla.ac.uk>
Date: Tue, 10 Sep 2024 19:41:10 +0100
Subject: [PATCH] test: update mocket get_loudnesses & test _keep_num

---
 test/test_feature_extractors.py | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/test/test_feature_extractors.py b/test/test_feature_extractors.py
index d19c93e..914e11c 100644
--- a/test/test_feature_extractors.py
+++ b/test/test_feature_extractors.py
@@ -220,15 +220,33 @@ class TestLoudAudioFeatureExtractor(unittest.TestCase):
         self.assertEqual(len(test_extractor.features), 5)
 
     def test_extract_mocked_get_loudnesses(self):
-        """Test extract with mocked loudness detection"""
+        """Test extract with mocked loudness detection - 100 loudnesses generated"""
         video_source = TestSourceMedia().one_colour_silent_audio()
-        test_extractor = extractors.LoudAudioFeatureExtractor(input_files=video_source)
+        test_extractor = extractors.LoudAudioFeatureExtractor(input_files=video_source, num_features=100)
         test_extractor._get_loudnesses = self._mock_get_loudnessess
         test_extractor.setup()
         test_extractor.run()
         test_extractor.teardown()
         self.assertEqual(len(test_extractor.features), 100)
 
+    def test_keep_num(self):
+        """Test keep_num correctly keeps 5 / 10"""
+        video_source = TestSourceMedia().one_colour_silent_audio()
+        with self.subTest("keep 5 (default)"):
+            test_extractor = extractors.LoudAudioFeatureExtractor(input_files=video_source)
+            test_extractor._get_loudnesses = self._mock_get_loudnessess
+            test_extractor.setup()
+            test_extractor.run()
+            test_extractor.teardown()
+            self.assertEqual(len(test_extractor.features), 5)
+        with self.subTest("keep 10"):
+            test_extractor = extractors.LoudAudioFeatureExtractor(input_files=video_source, num_features=10)
+            test_extractor._get_loudnesses = self._mock_get_loudnessess
+            test_extractor.setup()
+            test_extractor.run()
+            test_extractor.teardown()
+            self.assertEqual(len(test_extractor.features), 10)
+
     # TODO: add sample video with loud audio to test _loudnessdetect()
 
 class TestVideoActivityFeatureExtractor(unittest.TestCase):