Remove time inefficient loops in mlbf

Remove test

Remove all time inefficient loops from mlbf

.github/workflows/_test.yml

@@ -60,6 +60,11 @@ jobs:
             services: ''
             compose_file: docker-compose.yml:docker-compose.ci.yml
             run: make test_es_tests
+          -
+            name: MLBF
+            services: ''
+            compose_file: docker-compose.yml:docker-compose.ci.yml
+            run: make test_mlbf
           -
             name: Codestyle
             services: web

Makefile-docker

@@ -163,6 +163,13 @@ test_es: ## run the ES tests
 		-m es_tests \
 		$(ARGS)
 
+.PHONY: test_mlbf
+test_mlbf: ## run the MLBF tests
+	pytest \
+		$(PYTEST_SRC) \
+		-m mlbf_tests \
+		$(ARGS)
+
 .PHONY: test_no_es
 test_no_es: ## run all but the ES tests
 	pytest \

pyproject.toml

@@ -80,6 +80,9 @@ norecursedirs = [
   "venv",
   "__pycache__",
 ]
+# Disables the timeout globally. We can investigate
+# setting a SLA for tests in the future.
+timeout = 0
 DJANGO_SETTINGS_MODULE = "settings_test"
 # Ignoring csp deprecation warnings, we have control over the module and
 # currently it warns for child-src which is deprecated in CSPv3 but we're still

requirements/dev.txt

@@ -311,3 +311,6 @@ pytest-reportlog==0.4.0 \
 django-dbbackup==4.2.1 \
     --hash=sha256:157a2ec10d482345cd75092e510ac40d6e2ee6084604a1d17abe178c2f06bc69 \
     --hash=sha256:b23265600ead0780ca781b1b4b594949aaa8a20d74f08701f91ee9d7eb1f08cd
+pytest-timeout==2.3.1 \
+    --hash=sha256:12397729125c6ecbdaca01035b9e5239d4db97352320af155b3f5de1ba5165d9 \
+    --hash=sha256:68188cb703edfc6a18fad98dc25a3c61e9f24d644b0b70f33af545219fc7813e

src/olympia/blocklist/mlbf.py

@@ -2,7 +2,8 @@ import json
 import os
 import secrets
 from enum import Enum
-from typing import Dict, List, Optional, Tuple
+from functools import lru_cache
+from typing import List, Optional, Tuple
 
 from django.utils.functional import cached_property
 
@@ -21,15 +22,47 @@ log = olympia.core.logger.getLogger('z.amo.blocklist')
 
 
 def ordered_diff_lists(
-    previous: List[str], current: List[str]
+    current: List[str],
+    previous: List[str],
 ) -> Tuple[List[str], List[str], int]:
-    current_set = set(current)
     previous_set = set(previous)
     # Use lists instead of sets to maintain order
-    extras = [x for x in current if x not in previous_set]
-    deletes = [x for x in previous if x not in current_set]
-    changed_count = len(extras) + len(deletes)
-    return extras, deletes, changed_count
+    return [x for x in current if x not in previous_set]
+
+
+def get_not_blocked_items(all_blocked_version_ids: List[int]):
+    """
+    Returns a list of tuples containing the guid, version of all not blocked
+    versions. We use distinct to avoid duplicates, order by ID to ensure
+    cache.json is always sorted consistently, and return the values as a tuple
+    to make it easier to mock in tests.
+    """
+    return (
+        Version.unfiltered.exclude(id__in=all_blocked_version_ids or ())
+        .distinct()
+        .order_by('id')
+        .values_list('addon__addonguid__guid', 'version')
+    )
+
+
+def get_all_blocked_items():
+    """
+    Returns a list of tuples containing the guid, version, version_id, and
+    block_type of all blocked items. We use distinct to avoid duplicates,
+    Order by ID to ensure cache.json is always sorted consistently,
+    and return the values as a tuple to make it easier to mock in tests.
+    """
+    return (
+        BlockVersion.objects.filter(version__file__is_signed=True)
+        .distinct()
+        .order_by('id')
+        .values_list(
+            'block__guid',
+            'version__version',
+            'version_id',
+            'block_type',
+        )
+    )
 
 
 def generate_mlbf(stats, blocked, not_blocked):
@@ -136,50 +169,46 @@ class MLBFDataBaseLoader(BaseMLBFLoader):
 
     @cached_property
     def _all_blocks(self):
-        return (
-            BlockVersion.objects.filter(version__file__is_signed=True)
-            .distinct()
-            .order_by('id')
-            .values_list(
-                'block__guid',
-                'version__version',
-                'version_id',
-                'block_type',
-                named=True,
-            )
-        )
+        _blocked_version_ids = []
+        _blocked = {
+            BlockType.BLOCKED: [],
+            BlockType.SOFT_BLOCKED: [],
+        }
 
-    def _format_blocks(self, block_type: BlockType) -> List[str]:
-        return MLBF.hash_filter_inputs(
-            [
-                (version.block__guid, version.version__version)
-                for version in self._all_blocks
-                if version.block_type == block_type
-            ]
-        )
+        # We define get_all_blocked_items as a separate function to allow
+        # mocking the database query in tests to simulate large data sets.
+        for guid, version_string, version_id, block_type in get_all_blocked_items():
+            _blocked_version_ids.append(version_id)
+            _blocked[block_type].append((guid, version_string))
+
+        return _blocked, _blocked_version_ids
 
     @cached_property
-    def blocked_items(self) -> List[str]:
-        return self._format_blocks(BlockType.BLOCKED)
+    def blocked_items(self):
+        blocks_dict, _ = self._all_blocks
+        return MLBF.hash_filter_inputs(blocks_dict[BlockType.BLOCKED])
 
     @cached_property
-    def soft_blocked_items(self) -> List[str]:
-        return self._format_blocks(BlockType.SOFT_BLOCKED)
+    def soft_blocked_items(self):
+        blocks_dict, _ = self._all_blocks
+        return MLBF.hash_filter_inputs(blocks_dict[BlockType.SOFT_BLOCKED])
 
     @cached_property
-    def not_blocked_items(self) -> List[str]:
-        all_blocks_ids = [version.version_id for version in self._all_blocks]
+    def not_blocked_items(self):
+        _, all_blocked_version_ids = self._all_blocks
+        # We define not_blocked_items as a separate function to allow
+        # tests to simulate large data sets.
         not_blocked_items = MLBF.hash_filter_inputs(
-            Version.unfiltered.exclude(id__in=all_blocks_ids or ())
-            .distinct()
-            .order_by('id')
-            .values_list('addon__addonguid__guid', 'version')
+            get_not_blocked_items(all_blocked_version_ids)
         )
         blocked_items = set(self.blocked_items + self.soft_blocked_items)
         # even though we exclude all the version ids in the query there's an
         # edge case where the version string occurs twice for an addon so we
         # ensure not_blocked_items contain no blocked_items or soft_blocked_items.
-        return [item for item in not_blocked_items if item not in blocked_items]
+        return ordered_diff_lists(
+            not_blocked_items,
+            blocked_items,
+        )
 
 
 class MLBF:
@@ -200,8 +229,8 @@ class MLBF:
         self.data: BaseMLBFLoader = data_class(storage=self.storage)
 
     @classmethod
-    def hash_filter_inputs(cls, input_list):
-        """Returns a list"""
+    def hash_filter_inputs(cls, input_list: List[Tuple[str, str]]) -> List[str]:
+        """Returns a list of hashed strings"""
         return [
             cls.KEY_FORMAT.format(guid=guid, version=version)
             for (guid, version) in input_list
@@ -233,16 +262,17 @@ class MLBF:
 
         log.info(json.dumps(stats))
 
+    @lru_cache(maxsize=128)  # noqa: B019
     def generate_diffs(
-        self, previous_mlbf: 'MLBF' = None
-    ) -> Dict[BlockType, Tuple[List[str], List[str], int]]:
-        return {
-            block_type: ordered_diff_lists(
-                [] if previous_mlbf is None else previous_mlbf.data[block_type],
-                self.data[block_type],
-            )
-            for block_type in BlockType
-        }
+        self,
+        block_type: BlockType,
+        previous_mlbf: 'MLBF' = None,
+    ):
+        current = self.data[block_type]
+        previous = [] if previous_mlbf is None else previous_mlbf.data[block_type]
+        added = ordered_diff_lists(current, previous)
+        removed = ordered_diff_lists(previous, current)
+        return added, removed
 
     def generate_and_write_stash(self, previous_mlbf: 'MLBF' = None):
         """
@@ -269,21 +299,23 @@ class MLBF:
         in the "unblocked" stash (like for hard blocked items) as this would
         result in the version being in both blocked and unblocked stashes.
         """
-        diffs = self.generate_diffs(previous_mlbf)
-        blocked_added, blocked_removed, _ = diffs[BlockType.BLOCKED]
+        blocked_added, blocked_removed = self.generate_diffs(
+            BlockType.BLOCKED, previous_mlbf
+        )
         stash_json = {
             'blocked': blocked_added,
             'unblocked': blocked_removed,
         }
 
         if waffle.switch_is_active('enable-soft-blocking'):
-            soft_blocked_added, soft_blocked_removed, _ = diffs[BlockType.SOFT_BLOCKED]
+            soft_blocked_added, soft_blocked_removed = self.generate_diffs(
+                BlockType.SOFT_BLOCKED, previous_mlbf
+            )
             stash_json['softblocked'] = soft_blocked_added
-            stash_json['unblocked'] = [
-                unblocked
-                for unblocked in (blocked_removed + soft_blocked_removed)
-                if unblocked not in blocked_added
-            ]
+            stash_json['unblocked'] = ordered_diff_lists(
+                blocked_removed + soft_blocked_removed,
+                blocked_added,
+            )
 
         # write stash
         stash_path = self.stash_path
@@ -295,8 +327,8 @@ class MLBF:
     def blocks_changed_since_previous(
         self, block_type: BlockType = BlockType.BLOCKED, previous_mlbf: 'MLBF' = None
     ):
-        _, _, changed_count = self.generate_diffs(previous_mlbf)[block_type]
-        return changed_count
+        added, removed = self.generate_diffs(block_type, previous_mlbf)
+        return len(added) + len(removed)
 
     @classmethod
     def load_from_storage(

src/olympia/blocklist/tests/test_cron.py

@@ -34,6 +34,7 @@ from olympia.zadmin.models import set_config
 STATSD_PREFIX = 'blocklist.cron.upload_mlbf_to_remote_settings.'
 
 
+@pytest.mark.mlbf_tests
 @freeze_time('2020-01-01 12:34:56')
 @override_switch('blocklist_mlbf_submit', active=True)
 class TestUploadToRemoteSettings(TestCase):
@@ -379,6 +380,48 @@ class TestUploadToRemoteSettings(TestCase):
             in self.mocks['olympia.blocklist.cron.upload_filter.delay'].call_args_list
         )
 
+    # This test is very slow, so we increase the timeout to 180 seconds.
+    @pytest.mark.timeout(120)
+    @mock.patch('olympia.blocklist.mlbf.get_not_blocked_items')
+    @mock.patch('olympia.blocklist.mlbf.get_all_blocked_items')
+    def test_large_data_set(
+        self,
+        mock_get_all_blocked_items,
+        mock_get_not_blocked_items,
+    ):
+        """
+        Test that the cron can handle a large data set
+        We can safely mock the filter cascase initialize and verify methods
+        because they massively impact performance of the test and we don't
+        need to test them here.
+        """
+        two_million_blocked = [
+            (
+                f'{x}@blocked',
+                f'{x}@version',
+                x,
+                # even numbers are blocked, odd numbers are soft blocked
+                BlockType.BLOCKED if x % 2 == 0 else BlockType.SOFT_BLOCKED,
+            )
+            for x in range(0, 2_000_000)
+        ]
+        one_million_not_blocked = [
+            (f'{x}@not_blocked', f'{x}@version') for x in range(2_000_000, 3_000_000)
+        ]
+
+        mock_get_all_blocked_items.return_value = two_million_blocked
+        mock_get_not_blocked_items.return_value = one_million_not_blocked
+
+        upload_mlbf_to_remote_settings()
+
+        mlbf = MLBF.load_from_storage(self.current_time)
+        assert len(mlbf.data.blocked_items) == 1_000_000
+        assert len(mlbf.data.soft_blocked_items) == 1_000_000
+        assert len(mlbf.data.not_blocked_items) == 1_000_000
+
+        with override_switch('enable-soft-blocking', active=True):
+            upload_mlbf_to_remote_settings()
+
 
 class TestTimeMethods(TestCase):
     @freeze_time('2024-10-10 12:34:56')

src/olympia/blocklist/tests/test_mlbf.py

@@ -1,6 +1,7 @@
 import json
 from functools import cached_property
 
+import pytest
 from waffle.testutils import override_switch
 
 from olympia import amo
@@ -29,6 +30,7 @@ class _MLBFBase(TestCase):
     def setUp(self):
         self.storage = SafeStorage()
         self.user = user_factory()
+        self.empty_diff = ([], [])
 
     def _blocked_addon(self, block_type=BlockType.BLOCKED, **kwargs):
         addon = addon_factory(**kwargs)
@@ -50,13 +52,13 @@ class _MLBFBase(TestCase):
 
 class TestOrderedDiffLists(TestCase):
     def test_return_added(self):
-        assert ordered_diff_lists(['a', 'b'], ['a', 'b', 'c']) == (['c'], [], 1)
+        assert ordered_diff_lists(['a', 'b'], ['a', 'b', 'c']) == (['c'], [])
 
     def test_return_removed(self):
-        assert ordered_diff_lists(['a', 'b', 'c'], ['a', 'b']) == ([], ['c'], 1)
+        assert ordered_diff_lists(['a', 'b', 'c'], ['a', 'b']) == ([], ['c'])
 
     def test_return_added_and_removed(self):
-        assert ordered_diff_lists(['a', 'b', 'c'], ['b', 'c', 'd']) == (['d'], ['a'], 2)
+        assert ordered_diff_lists(['a', 'b', 'c'], ['b', 'c', 'd']) == (['d'], ['a'])
 
     def test_large_diff(self):
         size = 2_000_000
@@ -65,7 +67,6 @@ class TestOrderedDiffLists(TestCase):
         assert ordered_diff_lists(even_items, odd_items) == (
             odd_items,
             even_items,
-            size,
         )
 
 
@@ -261,6 +262,7 @@ class TestMLBFDataBaseLoader(_MLBFBase):
         assert default == ['guid:version']
 
 
+@pytest.mark.mlbf_tests
 class TestMLBF(_MLBFBase):
     def test_get_data_from_db(self):
         self._blocked_addon()
@@ -333,28 +335,28 @@ class TestMLBF(_MLBFBase):
         addon, _ = self._blocked_addon(file_kw={'is_signed': True})
         base_mlbf = MLBF.generate_from_db('base')
 
-        stateless_diff = {
-            BlockType.BLOCKED: (
-                MLBF.hash_filter_inputs(
-                    [(addon.block.guid, addon.current_version.version)]
-                ),
-                [],
-                1,
+        blocked_diff = (
+            MLBF.hash_filter_inputs(
+                [(addon.block.guid, addon.current_version.version)]
             ),
-            BlockType.SOFT_BLOCKED: ([], [], 0),
-        }
+            [],
+        )
 
-        assert base_mlbf.generate_diffs() == stateless_diff
+        assert base_mlbf.generate_diffs(BlockType.BLOCKED) == blocked_diff
+        assert base_mlbf.blocks_changed_since_previous(BlockType.BLOCKED) == 1
+        assert base_mlbf.generate_diffs(BlockType.SOFT_BLOCKED) == self.empty_diff
 
         next_mlbf = MLBF.generate_from_db('next')
         # If we don't include the base_mlbf, unblocked version will still be in the diff
-        assert next_mlbf.generate_diffs() == stateless_diff
+        assert next_mlbf.generate_diffs(BlockType.BLOCKED) == blocked_diff
+        assert next_mlbf.generate_diffs(BlockType.SOFT_BLOCKED) == self.empty_diff
         # Providing a previous mlbf with the unblocked version already included
         # removes it from the diff
-        assert next_mlbf.generate_diffs(previous_mlbf=base_mlbf) == {
-            BlockType.BLOCKED: ([], [], 0),
-            BlockType.SOFT_BLOCKED: ([], [], 0),
-        }
+        assert next_mlbf.generate_diffs(BlockType.BLOCKED, base_mlbf) == self.empty_diff
+        assert (
+            next_mlbf.generate_diffs(BlockType.SOFT_BLOCKED, base_mlbf)
+            == self.empty_diff
+        )
 
     def test_diff_no_changes(self):
         addon, block = self._blocked_addon()
@@ -362,10 +364,11 @@ class TestMLBF(_MLBFBase):
         base_mlbf = MLBF.generate_from_db('test')
         next_mlbf = MLBF.generate_from_db('test_two')
 
-        assert next_mlbf.generate_diffs(previous_mlbf=base_mlbf) == {
-            BlockType.BLOCKED: ([], [], 0),
-            BlockType.SOFT_BLOCKED: ([], [], 0),
-        }
+        assert next_mlbf.generate_diffs(BlockType.BLOCKED, base_mlbf) == self.empty_diff
+        assert (
+            next_mlbf.generate_diffs(BlockType.SOFT_BLOCKED, base_mlbf)
+            == self.empty_diff
+        )
 
     def test_diff_block_added(self):
         addon, block = self._blocked_addon()
@@ -377,16 +380,19 @@ class TestMLBF(_MLBFBase):
 
         next_mlbf = MLBF.generate_from_db('test_two')
 
-        assert next_mlbf.generate_diffs(previous_mlbf=base_mlbf) == {
-            BlockType.BLOCKED: (
-                MLBF.hash_filter_inputs(
-                    [(new_block.block.guid, new_block.version.version)]
-                ),
-                [],
-                1,
+        assert next_mlbf.generate_diffs(BlockType.BLOCKED, base_mlbf) == (
+            MLBF.hash_filter_inputs(
+                [(new_block.block.guid, new_block.version.version)]
             ),
-            BlockType.SOFT_BLOCKED: ([], [], 0),
-        }
+            [],
+        )
+        assert (
+            next_mlbf.blocks_changed_since_previous(BlockType.BLOCKED, base_mlbf) == 1
+        )
+        assert (
+            next_mlbf.generate_diffs(BlockType.SOFT_BLOCKED, base_mlbf)
+            == self.empty_diff
+        )
 
     def test_diff_block_removed(self):
         addon, block = self._blocked_addon()
@@ -397,16 +403,19 @@ class TestMLBF(_MLBFBase):
         block_version.delete()
         next_mlbf = MLBF.generate_from_db('test_two')
 
-        assert next_mlbf.generate_diffs(previous_mlbf=base_mlbf) == {
-            BlockType.BLOCKED: (
-                [],
-                MLBF.hash_filter_inputs(
-                    [(block_version.block.guid, block_version.version.version)]
-                ),
-                1,
+        assert next_mlbf.generate_diffs(BlockType.BLOCKED, base_mlbf) == (
+            [],
+            MLBF.hash_filter_inputs(
+                [(block_version.block.guid, block_version.version.version)]
             ),
-            BlockType.SOFT_BLOCKED: ([], [], 0),
-        }
+        )
+        assert (
+            next_mlbf.blocks_changed_since_previous(BlockType.BLOCKED, base_mlbf) == 1
+        )
+        assert (
+            next_mlbf.generate_diffs(BlockType.SOFT_BLOCKED, base_mlbf)
+            == self.empty_diff
+        )
 
     def test_diff_block_added_and_removed(self):
         addon, block = self._blocked_addon()
@@ -422,18 +431,21 @@ class TestMLBF(_MLBFBase):
 
         next_mlbf = MLBF.generate_from_db('test_two')
 
-        assert next_mlbf.generate_diffs(previous_mlbf=base_mlbf) == {
-            BlockType.BLOCKED: (
-                MLBF.hash_filter_inputs(
-                    [(new_block.block.guid, new_block.version.version)]
-                ),
-                MLBF.hash_filter_inputs(
-                    [(block_version.block.guid, block_version.version.version)]
-                ),
-                2,
+        assert next_mlbf.generate_diffs(BlockType.BLOCKED, base_mlbf) == (
+            MLBF.hash_filter_inputs(
+                [(new_block.block.guid, new_block.version.version)]
             ),
-            BlockType.SOFT_BLOCKED: ([], [], 0),
-        }
+            MLBF.hash_filter_inputs(
+                [(block_version.block.guid, block_version.version.version)]
+            ),
+        )
+        assert (
+            next_mlbf.blocks_changed_since_previous(BlockType.BLOCKED, base_mlbf) == 2
+        )
+        assert (
+            next_mlbf.generate_diffs(BlockType.SOFT_BLOCKED, base_mlbf)
+            == self.empty_diff
+        )
 
     def test_diff_block_hard_to_soft(self):
         addon, block = self._blocked_addon()
@@ -444,22 +456,25 @@ class TestMLBF(_MLBFBase):
         block_version.update(block_type=BlockType.SOFT_BLOCKED)
         next_mlbf = MLBF.generate_from_db('test_two')
 
-        assert next_mlbf.generate_diffs(previous_mlbf=base_mlbf) == {
-            BlockType.BLOCKED: (
-                [],
-                MLBF.hash_filter_inputs(
-                    [(block_version.block.guid, block_version.version.version)]
-                ),
-                1,
+        assert next_mlbf.generate_diffs(BlockType.BLOCKED, base_mlbf) == (
+            [],
+            MLBF.hash_filter_inputs(
+                [(block_version.block.guid, block_version.version.version)]
             ),
-            BlockType.SOFT_BLOCKED: (
-                MLBF.hash_filter_inputs(
-                    [(block_version.block.guid, block_version.version.version)]
-                ),
-                [],
-                1,
+        )
+        assert (
+            next_mlbf.blocks_changed_since_previous(BlockType.BLOCKED, base_mlbf) == 1
+        )
+        assert next_mlbf.generate_diffs(BlockType.SOFT_BLOCKED, base_mlbf) == (
+            MLBF.hash_filter_inputs(
+                [(block_version.block.guid, block_version.version.version)]
             ),
-        }
+            [],
+        )
+        assert (
+            next_mlbf.blocks_changed_since_previous(BlockType.SOFT_BLOCKED, base_mlbf)
+            == 1
+        )
 
     def test_diff_block_soft_to_hard(self):
         addon, block = self._blocked_addon()
@@ -470,22 +485,25 @@ class TestMLBF(_MLBFBase):
         block_version.update(block_type=BlockType.BLOCKED)
         next_mlbf = MLBF.generate_from_db('test_two')
 
-        assert next_mlbf.generate_diffs(previous_mlbf=base_mlbf) == {
-            BlockType.BLOCKED: (
-                MLBF.hash_filter_inputs(
-                    [(block_version.block.guid, block_version.version.version)]
-                ),
-                [],
-                1,
+        assert next_mlbf.generate_diffs(BlockType.BLOCKED, base_mlbf) == (
+            MLBF.hash_filter_inputs(
+                [(block_version.block.guid, block_version.version.version)]
             ),
-            BlockType.SOFT_BLOCKED: (
-                [],
-                MLBF.hash_filter_inputs(
-                    [(block_version.block.guid, block_version.version.version)]
-                ),
-                1,
+            [],
+        )
+        assert (
+            next_mlbf.blocks_changed_since_previous(BlockType.BLOCKED, base_mlbf) == 1
+        )
+        assert next_mlbf.generate_diffs(BlockType.SOFT_BLOCKED, base_mlbf) == (
+            [],
+            MLBF.hash_filter_inputs(
+                [(block_version.block.guid, block_version.version.version)]
             ),
-        }
+        )
+        assert (
+            next_mlbf.blocks_changed_since_previous(BlockType.SOFT_BLOCKED, base_mlbf)
+            == 1
+        )
 
     def test_diff_invalid_cache(self):
         addon, block = self._blocked_addon(file_kw={'is_signed': True})
@@ -507,16 +525,17 @@ class TestMLBF(_MLBFBase):
 
         # The diff should include the soft blocked version because it was removed
         # and should not include the blocked version because it was not changed
-        assert mlbf.generate_diffs(previous_mlbf=previous_mlbf) == {
-            BlockType.BLOCKED: ([], [], 0),
-            BlockType.SOFT_BLOCKED: (
-                MLBF.hash_filter_inputs(
-                    [(soft_blocked.block.guid, soft_blocked.version.version)]
-                ),
-                [],
-                1,
+        assert mlbf.generate_diffs(BlockType.BLOCKED, previous_mlbf) == self.empty_diff
+        assert mlbf.generate_diffs(BlockType.SOFT_BLOCKED, previous_mlbf) == (
+            MLBF.hash_filter_inputs(
+                [(soft_blocked.block.guid, soft_blocked.version.version)]
             ),
-        }
+            [],
+        )
+        assert (
+            mlbf.blocks_changed_since_previous(BlockType.SOFT_BLOCKED, previous_mlbf)
+            == 1
+        )
 
     def test_diff_all_possible_changes(self):
         """
@@ -570,10 +589,16 @@ class TestMLBF(_MLBFBase):
         first_mlbf = MLBF.generate_from_db('first')
 
         # We expect the 4 blocked versions to be in the diff, sorted by block type.
-        assert first_mlbf.generate_diffs() == {
-            BlockType.BLOCKED: ([five_hash, six_hash], [], 2),
-            BlockType.SOFT_BLOCKED: ([three_hash, four_hash], [], 2),
-        }
+        assert first_mlbf.generate_diffs(BlockType.BLOCKED) == (
+            [five_hash, six_hash],
+            [],
+        )
+        assert first_mlbf.blocks_changed_since_previous(BlockType.BLOCKED) == 2
+        assert first_mlbf.generate_diffs(BlockType.SOFT_BLOCKED) == (
+            [three_hash, four_hash],
+            [],
+        )
+        assert first_mlbf.blocks_changed_since_previous(BlockType.SOFT_BLOCKED) == 2
 
         # The first time we generate the stash, we expect 3 to 6 to be in the stash
         # as they have some kind of block applied.
@@ -614,19 +639,24 @@ class TestMLBF(_MLBFBase):
         # The order is based on the ID (i.e. creation time) of the block,
         # not the version so we expect two after three since two was
         # blocked after three.
-        assert second_mlbf.generate_diffs(previous_mlbf=first_mlbf) == {
-            BlockType.BLOCKED: (
-                [three_hash, two_hash],
-                [five_hash, six_hash],
-                4,
-            ),
-            # Same as above, one had a block created after five so it comes second.
-            BlockType.SOFT_BLOCKED: (
-                [five_hash, one_hash],
-                [three_hash, four_hash],
-                4,
-            ),
-        }
+        assert second_mlbf.generate_diffs(BlockType.BLOCKED, first_mlbf) == (
+            [three_hash, two_hash],
+            [five_hash, six_hash],
+        )
+        assert (
+            second_mlbf.blocks_changed_since_previous(BlockType.BLOCKED, first_mlbf)
+            == 4
+        )
+        assert second_mlbf.generate_diffs(BlockType.SOFT_BLOCKED, first_mlbf) == (
+            [five_hash, one_hash],
+            [three_hash, four_hash],
+        )
+        assert (
+            second_mlbf.blocks_changed_since_previous(
+                BlockType.SOFT_BLOCKED, first_mlbf
+            )
+            == 4
+        )
 
         assert second_mlbf.generate_and_write_stash(previous_mlbf=first_mlbf) == {
             'blocked': [three_hash, two_hash],

src/olympia/blocklist/tests/test_tasks.py

@@ -124,6 +124,7 @@ class TestProcessBlocklistSubmission(TransactionTestCase):
 
 
 @pytest.mark.django_db
+@pytest.mark.mlbf_tests
 class TestUploadMLBFToRemoteSettings(TestCase):
     def setUp(self):
         self.user = user_factory()