From 66263f03711c28e4ffbedcff8f4202200bc0da2d Mon Sep 17 00:00:00 2001
From: Gretchen Miller <gretchen@gretchenleigh.zone>
Date: Fri, 13 Sep 2024 14:21:59 -0700
Subject: [PATCH] WT-2950 video capture options enum

---
 brozzler/model.py  | 7 ++++++-
 brozzler/worker.py | 3 ++-
 brozzler/ydl.py    | 3 ++-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/brozzler/model.py b/brozzler/model.py
index 1603142..9766913 100644
--- a/brozzler/model.py
+++ b/brozzler/model.py
@@ -34,6 +34,7 @@ import urllib
 import uuid
 import yaml
 import zlib
+from enum import Enum
 from typing import Optional
 
 
@@ -219,6 +220,10 @@ class Job(doublethink.Document, ElapsedMixIn):
         self.status = "FINISHED"
         self.starts_and_stops[-1]["stop"] = doublethink.utcnow()
 
+class VideoCaptureOptions(Enum):
+    ENABLE_VIDEO_CAPTURE = "ENABLE_VIDEO_CAPTURE"
+    LIMIT_VIDEO_CAPTURE = "LIMIT_VIDEO_CAPTURE"
+    DISABLE_YTDLP_CAPTURE = "DISABLE_YTDLP_CAPTURE"
 
 class Site(doublethink.Document, ElapsedMixIn):
     logger = logging.getLogger(__module__ + "." + __qualname__)
@@ -236,7 +241,7 @@ class Site(doublethink.Document, ElapsedMixIn):
         if "scope" not in self:
             self.scope = {}
         if "video_capture" not in self:
-            self.video_capture = "ENABLE_VIDEO_CAPTURE"
+            self.video_capture = VideoCaptureOptions.ENABLE_VIDEO_CAPTURE.value
 
         # backward compatibility
         if "surt" in self.scope:
diff --git a/brozzler/worker.py b/brozzler/worker.py
index a0a4ea4..8601046 100644
--- a/brozzler/worker.py
+++ b/brozzler/worker.py
@@ -21,6 +21,7 @@ limitations under the License.
 import logging
 import brozzler
 import brozzler.browser
+from brozzler.model import VideoCaptureOptions
 import threading
 import time
 import urllib.request
@@ -250,7 +251,7 @@ class BrozzlerWorker:
             self.logger.info("needs fetch: %s", page)
             if (
                 self._is_video_type(page_headers)
-                and site.video_capture != "ENABLE_VIDEO_CAPTURE"
+                and site.video_capture != VideoCaptureOptions.ENABLE_VIDEO_CAPTURE.value
             ):
                 self.logger.info(
                     "skipping video content: video capture disabled for site"
diff --git a/brozzler/ydl.py b/brozzler/ydl.py
index a587d2e..c41a92b 100644
--- a/brozzler/ydl.py
+++ b/brozzler/ydl.py
@@ -20,6 +20,7 @@ import logging
 import yt_dlp
 from yt_dlp.utils import match_filter_func
 import brozzler
+from brozzler.model import VideoCaptureOptions
 import urllib.request
 import tempfile
 import urlcanon
@@ -38,7 +39,7 @@ def should_ytdlp(site, page, page_status):
     if page_status != 200:
         logging.info("skipping ytdlp: non-200 page status %s", page_status)
         return False
-    if site.video_capture != "ENABLE_VIDEO_CAPTURE":
+    if site.video_capture != VideoCaptureOptions.ENABLE_VIDEO_CAPTURE.value:
         logging.info("skipping ytdlp: site has video capture disabled")
         return False