Add tests

apache · Dec 22, 2021 · a1ef4fd · a1ef4fd
1 parent fdfa801
commit a1ef4fd
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 44 deletions.
diff --git a/tests/integration_tests/reports/commands_tests.py b/tests/integration_tests/reports/commands_tests.py
@@ -135,6 +135,7 @@ def create_report_notification(
     grace_period: Optional[int] = None,
     report_format: Optional[ReportDataFormat] = None,
     name: Optional[str] = None,
+    force_screenshot: bool = False,
 ) -> ReportSchedule:
     report_type = report_type or ReportScheduleType.REPORT
     target = email_target or slack_channel
@@ -174,6 +175,7 @@ def create_report_notification(
         validator_config_json=validator_config_json,
         grace_period=grace_period,
         report_format=report_format or ReportDataFormat.VISUALIZATION,
+        force_screenshot=force_screenshot,
     )
     return report_schedule
 
@@ -218,6 +220,18 @@ def create_report_email_chart():
         cleanup_report_schedule(report_schedule)
 
 
+@pytest.fixture()
+def create_report_email_chart_force_screenshot():
+    with app.app_context():
+        chart = db.session.query(Slice).first()
+        report_schedule = create_report_notification(
+            email_target="[email protected]", chart=chart, force_screenshot=True
+        )
+        yield report_schedule
+
+        cleanup_report_schedule(report_schedule)
+
+
 @pytest.fixture()
 def create_report_email_chart_with_csv():
     with app.app_context():
@@ -368,9 +382,7 @@ def create_alert_slack_chart_success():
 
 
 @pytest.fixture(
-    params=[
-        "alert1",
-    ]
+    params=["alert1",]
 )
 def create_alert_slack_chart_grace(request):
     param_config = {
@@ -482,6 +494,7 @@ def create_alert_email_chart(request):
                 validator_config_json=param_config[request.param][
                     "validator_config_json"
                 ],
+                force_screenshot=True,
             )
             yield report_schedule
 
@@ -647,9 +660,7 @@ def create_invalid_sql_alert_email_chart(request):
 @patch("superset.reports.notifications.email.send_email_smtp")
 @patch("superset.utils.screenshots.ChartScreenshot.get_screenshot")
 def test_email_chart_report_schedule(
-    screenshot_mock,
-    email_mock,
-    create_report_email_chart,
+    screenshot_mock, email_mock, create_report_email_chart,
 ):
     """
     ExecuteReport Command: Test chart email report schedule with screenshot
@@ -682,15 +693,56 @@ def test_email_chart_report_schedule(
         assert_log(ReportState.SUCCESS)
 
 
+@pytest.mark.usefixtures(
+    "load_birth_names_dashboard_with_slices",
+    "create_report_email_chart_force_screenshot",
+)
+@patch("superset.reports.notifications.email.send_email_smtp")
+@patch("superset.utils.screenshots.ChartScreenshot.get_screenshot")
+def test_email_chart_report_schedule_force_screenshot(
+    screenshot_mock, email_mock, create_report_email_chart_force_screenshot,
+):
+    """
+    ExecuteReport Command: Test chart email report schedule with screenshot
+
+    In this test ``force_screenshot`` is true, and the screenshot URL should
+    reflect that.
+    """
+    # setup screenshot mock
+    screenshot_mock.return_value = SCREENSHOT_FILE
+
+    with freeze_time("2020-01-01T00:00:00Z"):
+        AsyncExecuteReportScheduleCommand(
+            TEST_ID, create_report_email_chart_force_screenshot.id, datetime.utcnow()
+        ).run()
+
+        notification_targets = get_target_from_report_schedule(
+            create_report_email_chart_force_screenshot
+        )
+        # assert that the link sent is correct
+        assert (
+            '<a href="http://0.0.0.0:8080/superset/explore/?'
+            "form_data=%7B%22slice_id%22%3A+"
+            f"{create_report_email_chart_force_screenshot.chart.id}%7D&"
+            'standalone=true&force=true">Explore in Superset</a>'
+            in email_mock.call_args[0][2]
+        )
+        # Assert the email smtp address
+        assert email_mock.call_args[0][0] == notification_targets[0]
+        # Assert the email inline screenshot
+        smtp_images = email_mock.call_args[1]["images"]
+        assert smtp_images[list(smtp_images.keys())[0]] == SCREENSHOT_FILE
+        # Assert logs are correct
+        assert_log(ReportState.SUCCESS)
+
+
 @pytest.mark.usefixtures(
     "load_birth_names_dashboard_with_slices", "create_alert_email_chart"
 )
 @patch("superset.reports.notifications.email.send_email_smtp")
 @patch("superset.utils.screenshots.ChartScreenshot.get_screenshot")
 def test_email_chart_alert_schedule(
-    screenshot_mock,
-    email_mock,
-    create_alert_email_chart,
+    screenshot_mock, email_mock, create_alert_email_chart,
 ):
     """
     ExecuteReport Command: Test chart email alert schedule with screenshot
@@ -764,9 +816,7 @@ def test_email_chart_alert_schedule(
 @patch("superset.reports.notifications.email.send_email_smtp")
 @patch("superset.utils.screenshots.ChartScreenshot.get_screenshot")
 def test_email_chart_report_dry_run(
-    screenshot_mock,
-    email_mock,
-    create_report_email_chart,
+    screenshot_mock, email_mock, create_report_email_chart,
 ):
     """
     ExecuteReport Command: Test chart email report schedule dry run
@@ -791,11 +841,7 @@ def test_email_chart_report_dry_run(
 @patch("superset.reports.notifications.email.send_email_smtp")
 @patch("superset.utils.csv.get_chart_csv_data")
 def test_email_chart_report_schedule_with_csv(
-    csv_mock,
-    email_mock,
-    mock_open,
-    mock_urlopen,
-    create_report_email_chart_with_csv,
+    csv_mock, email_mock, mock_open, mock_urlopen, create_report_email_chart_with_csv,
 ):
     """
     ExecuteReport Command: Test chart email report schedule with CSV
@@ -984,9 +1030,7 @@ def test_email_dashboard_report_schedule(
 @patch("superset.reports.notifications.slack.WebClient.files_upload")
 @patch("superset.utils.screenshots.ChartScreenshot.get_screenshot")
 def test_slack_chart_report_schedule(
-    screenshot_mock,
-    file_upload_mock,
-    create_report_slack_chart,
+    screenshot_mock, file_upload_mock, create_report_slack_chart,
 ):
     """
     ExecuteReport Command: Test chart slack report schedule
@@ -1205,9 +1249,7 @@ def test_report_schedule_success_grace_end(
 @patch("superset.reports.notifications.email.send_email_smtp")
 @patch("superset.utils.screenshots.ChartScreenshot.get_screenshot")
 def test_alert_limit_is_applied(
-    screenshot_mock,
-    email_mock,
-    create_alert_email_chart,
+    screenshot_mock, email_mock, create_alert_email_chart,
 ):
     """
     ExecuteReport Command: Test that all alerts apply a SQL limit to stmts
@@ -1263,9 +1305,7 @@ def test_email_dashboard_report_fails(
     ALERTS_ATTACH_REPORTS=True,
 )
 def test_slack_chart_alert(
-    screenshot_mock,
-    email_mock,
-    create_alert_email_chart,
+    screenshot_mock, email_mock, create_alert_email_chart,
 ):
     """
     ExecuteReport Command: Test chart slack alert
@@ -1322,9 +1362,7 @@ def test_slack_chart_alert_no_attachment(email_mock, create_alert_email_chart):
 @patch("superset.reports.notifications.slack.WebClient")
 @patch("superset.utils.screenshots.ChartScreenshot.get_screenshot")
 def test_slack_token_callable_chart_report(
-    screenshot_mock,
-    slack_client_mock_class,
-    create_report_slack_chart,
+    screenshot_mock, slack_client_mock_class, create_report_slack_chart,
 ):
     """
     ExecuteReport Command: Test chart slack alert (slack token callable)
@@ -1436,11 +1474,7 @@ def test_soft_timeout_screenshot(screenshot_mock, email_mock, create_alert_email
 @patch("superset.reports.notifications.email.send_email_smtp")
 @patch("superset.utils.csv.get_chart_csv_data")
 def test_soft_timeout_csv(
-    csv_mock,
-    email_mock,
-    mock_open,
-    mock_urlopen,
-    create_report_email_chart_with_csv,
+    csv_mock, email_mock, mock_open, mock_urlopen, create_report_email_chart_with_csv,
 ):
     """
     ExecuteReport Command: Test fail on generating csv
@@ -1464,8 +1498,7 @@ def test_soft_timeout_csv(
     assert email_mock.call_args[0][0] == OWNER_EMAIL
 
     assert_log(
-        ReportState.ERROR,
-        error_message="A timeout occurred while generating a csv.",
+        ReportState.ERROR, error_message="A timeout occurred while generating a csv.",
     )
 
 
@@ -1477,11 +1510,7 @@ def test_soft_timeout_csv(
 @patch("superset.reports.notifications.email.send_email_smtp")
 @patch("superset.utils.csv.get_chart_csv_data")
 def test_generate_no_csv(
-    csv_mock,
-    email_mock,
-    mock_open,
-    mock_urlopen,
-    create_report_email_chart_with_csv,
+    csv_mock, email_mock, mock_open, mock_urlopen, create_report_email_chart_with_csv,
 ):
     """
     ExecuteReport Command: Test fail on generating csv
@@ -1666,9 +1695,7 @@ def test_grace_period_error(email_mock, create_invalid_sql_alert_email_chart):
 @patch("superset.reports.notifications.email.send_email_smtp")
 @patch("superset.utils.screenshots.ChartScreenshot.get_screenshot")
 def test_grace_period_error_flap(
-    screenshot_mock,
-    email_mock,
-    create_invalid_sql_alert_email_chart,
+    screenshot_mock, email_mock, create_invalid_sql_alert_email_chart,
 ):
     """
     ExecuteReport Command: Test alert grace period on error

diff --git a/tests/integration_tests/reports/utils.py b/tests/integration_tests/reports/utils.py
@@ -51,6 +51,7 @@ def insert_report_schedule(
     recipients: Optional[List[ReportRecipients]] = None,
     report_format: Optional[ReportDataFormat] = None,
     logs: Optional[List[ReportExecutionLog]] = None,
+    force_screenshot: bool = False,
 ) -> ReportSchedule:
     owners = owners or []
     recipients = recipients or []
@@ -75,6 +76,7 @@ def insert_report_schedule(
         logs=logs,
         last_state=last_state,
         report_format=report_format,
+        force_screenshot=force_screenshot,
     )
     db.session.add(report_schedule)
     db.session.commit()