From 2d3959c479c2b25b097c8a9d3ea67cbb0ccfcf3a Mon Sep 17 00:00:00 2001
From: keisuke-umezawa <keisuke.umezawa@gmail.com>
Date: Sun, 11 Jun 2023 22:36:03 +0900
Subject: [PATCH] Migrate visual_regression_test

---
 e2e_tests/conftest.py               | 30 +++++++++++---------
 e2e_tests/visual_regression_test.py | 44 ++++++++++++++++++-----------
 2 files changed, 44 insertions(+), 30 deletions(-)

diff --git a/e2e_tests/conftest.py b/e2e_tests/conftest.py
index 58338334..07124540 100644
--- a/e2e_tests/conftest.py
+++ b/e2e_tests/conftest.py
@@ -1,11 +1,17 @@
-import pytest
-from playwright.sync_api import sync_playwright
-
 import threading
+from typing import Generator
 from wsgiref.simple_server import make_server
 
 import optuna
 from optuna_dashboard import wsgi
+from playwright.sync_api import Browser
+from playwright.sync_api import sync_playwright
+import pytest
+
+
+@pytest.fixture(scope="session")
+def port() -> int:
+    return 8081
 
 
 @pytest.fixture(scope="session")
@@ -180,18 +186,16 @@ def dummy_storage() -> optuna.storages.InMemoryStorage:
     return storage
 
 
-host = "127.0.0.1"
-port = 8081
-
-
-@pytest.fixture(scope='session')
-def server(request, dummy_storage):
+@pytest.fixture(scope="session")
+def server(
+    request: pytest.FixtureRequest, dummy_storage: optuna.storages.InMemoryStorage, port: int
+) -> None:
     app = wsgi(dummy_storage)
-    httpd = make_server(host, port, app)
+    httpd = make_server("127.0.0.1", port, app)
     thread = threading.Thread(target=httpd.serve_forever)
     thread.start()
 
-    def stop_server():
+    def stop_server() -> None:
         httpd.shutdown()
         httpd.server_close()
         thread.join()
@@ -199,8 +203,8 @@ def server(request, dummy_storage):
     request.addfinalizer(stop_server)
 
 
-@pytest.fixture(scope='module')
-def browser():
+@pytest.fixture(scope="module")
+def browser() -> Generator[Browser, None, None]:
     with sync_playwright() as playwright:
         browser = playwright.chromium.launch()
         yield browser
diff --git a/e2e_tests/visual_regression_test.py b/e2e_tests/visual_regression_test.py
index 8d8c108b..17135201 100644
--- a/e2e_tests/visual_regression_test.py
+++ b/e2e_tests/visual_regression_test.py
@@ -1,25 +1,35 @@
-import pytest
-
-from playwright.sync_api import expect
-from playwright.sync_api import Page
+import os
 
 import optuna
+from playwright.sync_api import Browser
 
 
-def test_first(
-    dummy_storage: optuna.storages.InMemoryStorage,
-    server,
-    browser,
+output_dir = "tmp"
+width = 1000
+height = 3000
+
+
+def test_take_screenshots_study_list(
+    server: None,
+    browser: Browser,
+    port: int,
 ) -> None:
-    print(dummy_storage)
-    print(id(dummy_storage))
+    url = f"http://localhost:{port}/dashboard/"
+    page = browser.new_page(viewport={"width": width, "height": height})
+    page.goto(url)
+    page.screenshot(path=os.path.join(output_dir, "study-list.png"))
 
-def test_second(
+
+def test_take_screenshots_studies(
     dummy_storage: optuna.storages.InMemoryStorage,
-    server,
-    browser,
+    server: None,
+    browser: Browser,
+    port: int,
 ) -> None:
-    page = browser.new_page()
-    page.goto("http://localhost:8081")
-    print(page.title())
-
+    summaries = optuna.get_all_study_summaries(dummy_storage)
+    study_ids = {s._study_id: s.study_name for s in summaries}
+    for study_id, study_name in study_ids.items():
+        url = f"http://localhost:{port}/dashboard/studies/{study_id}"
+        page = browser.new_page(viewport={"width": width, "height": height})
+        page.goto(url)
+        page.screenshot(path=os.path.join(output_dir, f"study-{study_name}.png"))