From ff3bab33d1874cc025c1dafbbc38c2d3bf4330fc Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Thu, 6 Feb 2025 09:10:02 -0500 Subject: [PATCH] Fix lint --- superset/utils/screenshots.py | 5 +++-- tests/unit_tests/utils/screenshot_test.py | 11 ++++++----- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/superset/utils/screenshots.py b/superset/utils/screenshots.py index 6abc03eda6451..7f40c0a8205f6 100644 --- a/superset/utils/screenshots.py +++ b/superset/utils/screenshots.py @@ -20,7 +20,7 @@ from datetime import datetime from enum import Enum from io import BytesIO -from typing import TYPE_CHECKING, TypedDict +from typing import cast, TYPE_CHECKING, TypedDict from flask import current_app @@ -73,7 +73,7 @@ class ScreenshotCachePayload: def __init__( self, image: bytes | None = None, - status: str = StatusValues.PENDING, + status: StatusValues = StatusValues.PENDING, timestamp: str = "", ): self._image = image @@ -209,6 +209,7 @@ def get_from_cache_key(cls, cache_key: str) -> ScreenshotCachePayload | None: elif isinstance(payload, ScreenshotCachePayload): pass elif isinstance(payload, dict): + payload = cast(ScreenshotCachePayloadType, payload) payload = ScreenshotCachePayload.from_dict(payload) return payload logger.info("Failed at getting from cache: %s", cache_key) diff --git a/tests/unit_tests/utils/screenshot_test.py b/tests/unit_tests/utils/screenshot_test.py index 1af00a36366cd..fa928fc73fb67 100644 --- a/tests/unit_tests/utils/screenshot_test.py +++ b/tests/unit_tests/utils/screenshot_test.py @@ -26,6 +26,7 @@ from superset.utils.screenshots import ( BaseScreenshot, ScreenshotCachePayload, + ScreenshotCachePayloadType, ) BASE_SCREENSHOT_PATH = "superset.utils.screenshots.BaseScreenshot" @@ -120,7 +121,7 @@ def _setup_compute_and_cache(self, mocker: MockerFixture, screenshot_obj): def test_happy_path(self, mocker: MockerFixture, screenshot_obj): self._setup_compute_and_cache(mocker, screenshot_obj) screenshot_obj.compute_and_cache(force=False) - cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") + cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key") assert cache_payload["status"] == "Updated" def test_screenshot_error(self, mocker: MockerFixture, screenshot_obj): @@ -128,7 +129,7 @@ def test_screenshot_error(self, mocker: MockerFixture, screenshot_obj): get_screenshot: MagicMock = mocks.get("get_screenshot") get_screenshot.side_effect = Exception screenshot_obj.compute_and_cache(force=False) - cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") + cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key") assert cache_payload["status"] == "Error" def test_resize_error(self, mocker: MockerFixture, screenshot_obj): @@ -136,7 +137,7 @@ def test_resize_error(self, mocker: MockerFixture, screenshot_obj): resize_image: MagicMock = mocks.get("resize_image") resize_image.side_effect = Exception screenshot_obj.compute_and_cache(force=False) - cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") + cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key") assert cache_payload["status"] == "Error" def test_skips_if_computing(self, mocker: MockerFixture, screenshot_obj): @@ -154,7 +155,7 @@ def test_skips_if_computing(self, mocker: MockerFixture, screenshot_obj): # Ensure that it processes when force = True screenshot_obj.compute_and_cache(force=True) get_screenshot.assert_called_once() - cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") + cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key") assert cache_payload["status"] == "Updated" def test_skips_if_updated(self, mocker: MockerFixture, screenshot_obj): @@ -176,7 +177,7 @@ def test_skips_if_updated(self, mocker: MockerFixture, screenshot_obj): force=True, window_size=window_size, thumb_size=thumb_size ) get_screenshot.assert_called_once() - cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") + cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key") assert cache_payload["image"] != b"initial_value" def test_resize(self, mocker: MockerFixture, screenshot_obj):