Support uploading stash/filter from multiple block types

src/olympia/blocklist/tests/test_tasks.py

@@ -1,6 +1,7 @@
 import json
 import os
 from datetime import datetime, timedelta
+from typing import Dict, List
 from unittest import TestCase, mock
 
 from django.conf import settings
@@ -432,6 +433,100 @@ def test_ignore_soft_blocked_if_switch_is_disabled(self):
                 self.generation_time, filter_list=[BlockType.SOFT_BLOCKED.name]
             )
             assert self.mocks['delete_record'].called
+    def test_upload_soft_blocked_filter(self):
+        self._test_upload_base_filter(BlockType.SOFT_BLOCKED)
+
+    def test_upload_soft_and_blocked_filter(self):
+        self._test_upload_base_filter(BlockType.BLOCKED, BlockType.SOFT_BLOCKED)
+
+    def _test_cleanup_old_records(
+        self,
+        filter_list: Dict[BlockType, int],
+        records: List[Dict[str, int]],
+        expected_calls: List[any],
+    ):
+        self._block_version(is_signed=True)
+        mlbf = MLBF.generate_from_db(self.generation_time)
+        for block_type, base_id in filter_list.items():
+            mlbf.generate_and_write_filter(block_type)
+            # We currently write to the old singular config key for hard blocks
+            # to preserve backward compatibility.
+            # In https://github.com/mozilla/addons/issues/15193
+            # we can remove this and start writing to the new plural key.
+            set_config(
+                MLBF_BASE_ID_CONFIG_KEY(block_type, compat=True),
+                base_id,
+                json_value=True,
+            )
+
+        self.mocks['records'].return_value = records
+        upload_filter.delay(
+            self.generation_time,
+            filter_list=[block_type.name for block_type in filter_list],
+        )
+
+        assert self.mocks['delete_record'].call_args_list == expected_calls
+
+        if len(filter_list.values()) > 0:
+            self.mocks['cleanup_old_files.delay'].assert_called_with(
+                base_filter_id=min(filter_list.values())
+            )
+            self.mocks['statsd.incr'].assert_called_with(
+                'blocklist.tasks.upload_filter.reset_collection'
+            )
+
+    def test_skip_cleanup_when_no_filters(self):
+        self._test_cleanup_old_records(
+            filter_list={},
+            records=[{'id': '0', 'generation_time': self.generation_time}],
+            expected_calls=[],
+        )
+
+    def test_cleanup_old_records(self):
+        """
+        Clean up 0 because its the only record matching the uplaoded filters
+        attachment_type or is older than genreation_time
+        """
+        self._test_cleanup_old_records(
+            filter_list={
+                BlockType.BLOCKED: self.generation_time,
+            },
+            records=[
+                self._attachment(0, 'bloomfilter-base', self.generation_time - 1),
+                self._attachment(
+                    1, 'softblocks-bloomfilter-base', self.generation_time
+                ),
+                self._stash(2, self.generation_time + 1),
+            ],
+            expected_calls=[mock.call(0)],
+        )
+
+    def test_cleanup_oldest_stash_records(self):
+        """
+        The oldest base id time is (self.generation_time -3)
+        Delete 0 as matching attachment
+        Delete 1 as older than base id
+        """
+        self._test_cleanup_old_records(
+            filter_list={
+                BlockType.BLOCKED: self.generation_time - 3,
+            },
+            records=[
+                # Deleted, matching attachment
+                self._attachment(0, 'bloomfilter-base', self.generation_time - 5),
+                self._stash(
+                    1, self.generation_time - 4
+                ),  # deleted older than oldest filter
+                # Not deleted, not matching attachment
+                self._attachment(
+                    1, 'softblocks-bloomfilter-base', self.generation_time - 3
+                ),
+                # Both Not delted, not older than oldest filter
+                self._stash(1, self.generation_time - 2),
+                self._stash(2, self.generation_time - 1),
+            ],
+            expected_calls=[mock.call(0), mock.call(1)],
+        )
 
     def test_create_stashed_filter(self):
         old_mlbf = MLBF.generate_from_db(self.generation_time - 1)
@@ -471,11 +566,13 @@ def test_create_stashed_filter(self):
 
     def test_raises_when_no_filter_exists(self):
         with self.assertRaises(FileNotFoundError):
-            upload_filter(self.generation_time, filter_list=[BlockType.BLOCKED.name])
+            upload_filter.delay(
+                self.generation_time, filter_list=[BlockType.BLOCKED.name]
+            )
 
     def test_raises_when_no_stash_exists(self):
         with self.assertRaises(FileNotFoundError):
-            upload_filter(self.generation_time, create_stash=True)
+            upload_filter.delay(self.generation_time, create_stash=True)
 
     def test_default_is_no_op(self):
         MLBF.generate_from_db(self.generation_time).generate_and_write_filter(