type and test fixes

archivebox/workers/tests/test_scheduled_crawls.py

@@ -1,7 +1,9 @@
 from datetime import timedelta
+from typing import cast
 from unittest.mock import patch
 
 from django.contrib.auth import get_user_model
+from django.contrib.auth.models import UserManager
 from django.test import TestCase
 from django.utils import timezone
 
@@ -12,7 +14,8 @@ from archivebox.workers.worker import CrawlWorker
 
 class TestScheduledCrawlMaterialization(TestCase):
     def setUp(self):
-        self.user = get_user_model().objects.create_user(
+        user_manager = cast(UserManager, get_user_model().objects)
+        self.user = user_manager.create_user(
             username='schedule-user',
             password='password',
         )
@@ -52,6 +55,8 @@ class TestScheduledCrawlMaterialization(TestCase):
         self.assertEqual(scheduled_crawls.count(), 2)
 
         queued_crawl = scheduled_crawls.last()
+        self.assertIsNotNone(queued_crawl)
+        assert queued_crawl is not None
         self.assertEqual(queued_crawl.status, Crawl.StatusChoices.QUEUED)
         self.assertEqual(queued_crawl.urls, 'https://example.com/feed.xml')
         self.assertEqual(queued_crawl.max_depth, 1)
@@ -63,7 +68,7 @@ class TestScheduledCrawlMaterialization(TestCase):
         Orchestrator(exit_on_idle=True)._materialize_due_schedules()
         self.assertEqual(Crawl.objects.filter(schedule=schedule).count(), 1)
 
-        Orchestrator(exit_on_idle=False, crawl_id=str(schedule.template_id))._materialize_due_schedules()
+        Orchestrator(exit_on_idle=False, crawl_id=str(schedule.template.id))._materialize_due_schedules()
         self.assertEqual(Crawl.objects.filter(schedule=schedule).count(), 1)
 
     @patch.object(CrawlWorker, 'start')

archivebox/workers/tests/test_snapshot_worker.py

@@ -1,5 +1,6 @@
 from pathlib import Path
 from types import SimpleNamespace
+from typing import Any, cast
 from unittest.mock import patch
 
 from django.test import SimpleTestCase
@@ -11,14 +12,14 @@ class TestSnapshotWorkerRetryForegroundHooks(SimpleTestCase):
     def _make_worker(self):
         worker = SnapshotWorker.__new__(SnapshotWorker)
         worker.pid = 12345
-        worker.snapshot = SimpleNamespace(
+        cast(Any, worker).snapshot = SimpleNamespace(
             status='started',
             refresh_from_db=lambda: None,
         )
         worker._snapshot_exceeded_hard_timeout = lambda: False
         worker._seal_snapshot_due_to_timeout = lambda: None
         worker._run_hook = lambda *args, **kwargs: SimpleNamespace()
-        worker._wait_for_hook = lambda *args, **kwargs: None
+        worker._wait_for_hook = lambda process, ar: None
         return worker
 
     @patch('archivebox.workers.worker.log_worker_event')
@@ -49,10 +50,10 @@ class TestSnapshotWorkerRetryForegroundHooks(SimpleTestCase):
             run_calls.append((args, kwargs))
             return SimpleNamespace()
 
-        def wait_for_hook(process, archive_result):
-            wait_calls.append((process, archive_result))
-            archive_result.status = 'succeeded'
-            archive_result.output_files = {'singlefile.html': {}}
+        def wait_for_hook(process, ar):
+            wait_calls.append((process, ar))
+            ar.status = 'succeeded'
+            ar.output_files = {'singlefile.html': {}}
 
         archive_result = SimpleNamespace(
             status='failed',