Skip to content

Commit

Permalink
[REF-723+] Upload with progress and cancellation (#1899)
Browse files Browse the repository at this point in the history
  • Loading branch information
@masenf
masenf authored and Alek Petuskey committed Nov 20, 2023
1 parent c74696c commit 4721c4f
Show file tree
Hide file tree
Showing 14 changed files with 646 additions and 180 deletions.
146 changes: 135 additions & 11 deletions integration/test_upload.py
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,14 @@
"""Integration tests for file upload."""
from __future__ import annotations

import asyncio
import time
from typing import Generator

import pytest
from selenium.webdriver.common.by import By

from reflex.testing import AppHarness
from reflex.testing import AppHarness, WebDriver


def UploadFile():
Expand All @@ -16,19 +17,36 @@ def UploadFile():

class UploadState(rx.State):
_file_data: dict[str, str] = {}
event_order: list[str] = []
progress_dicts: list[dict] = []

async def handle_upload(self, files: list[rx.UploadFile]):
for file in files:
upload_data = await file.read()
self._file_data[file.filename or ""] = upload_data.decode("utf-8")

async def handle_upload_secondary(self, files: list[rx.UploadFile]):
for file in files:
upload_data = await file.read()
self._file_data[file.filename or ""] = upload_data.decode("utf-8")
yield UploadState.chain_event

def upload_progress(self, progress):
assert progress
self.event_order.append("upload_progress")
self.progress_dicts.append(progress)

def chain_event(self):
self.event_order.append("chain_event")

def index():
return rx.vstack(
rx.input(
value=UploadState.router.session.client_token,
is_read_only=True,
id="token",
),
rx.heading("Default Upload"),
rx.upload(
rx.vstack(
rx.button("Select File"),
Expand All @@ -52,6 +70,47 @@ def index():
on_click=rx.clear_selected_files,
id="clear_button",
),
rx.heading("Secondary Upload"),
rx.upload(
rx.vstack(
rx.button("Select File"),
rx.text("Drag and drop files here or click to select files"),
),
id="secondary",
),
rx.button(
"Upload",
on_click=UploadState.handle_upload_secondary( # type: ignore
rx.upload_files(
upload_id="secondary",
on_upload_progress=UploadState.upload_progress,
),
),
id="upload_button_secondary",
),
rx.box(
rx.foreach(
rx.selected_files("secondary"),
lambda f: rx.text(f),
),
id="selected_files_secondary",
),
rx.button(
"Clear",
on_click=rx.clear_selected_files("secondary"),
id="clear_button_secondary",
),
rx.vstack(
rx.foreach(
UploadState.progress_dicts, # type: ignore
lambda d: rx.text(d.to_string()),
)
),
rx.button(
"Cancel",
on_click=rx.cancel_upload("secondary"),
id="cancel_button_secondary",
),
)

app = rx.App(state=UploadState)
Expand Down Expand Up @@ -94,14 +153,18 @@ def driver(upload_file: AppHarness):
driver.quit()


@pytest.mark.parametrize("secondary", [False, True])
@pytest.mark.asyncio
async def test_upload_file(tmp_path, upload_file: AppHarness, driver):
async def test_upload_file(
tmp_path, upload_file: AppHarness, driver: WebDriver, secondary: bool
):
"""Submit a file upload and check that it arrived on the backend.
Args:
tmp_path: pytest tmp_path fixture
upload_file: harness for UploadFile app.
driver: WebDriver instance.
secondary: whether to use the secondary upload form
"""
assert upload_file.app_instance is not None
token_input = driver.find_element(By.ID, "token")
Expand All @@ -110,9 +173,13 @@ async def test_upload_file(tmp_path, upload_file: AppHarness, driver):
token = upload_file.poll_for_value(token_input)
assert token is not None

upload_box = driver.find_element(By.XPATH, "//input[@type='file']")
suffix = "_secondary" if secondary else ""

upload_box = driver.find_elements(By.XPATH, "//input[@type='file']")[
1 if secondary else 0
]
assert upload_box
upload_button = driver.find_element(By.ID, "upload_button")
upload_button = driver.find_element(By.ID, f"upload_button{suffix}")
assert upload_button

exp_name = "test.txt"
Expand All @@ -132,9 +199,15 @@ async def get_file_data():
assert file_data[exp_name] == exp_contents

# check that the selected files are displayed
selected_files = driver.find_element(By.ID, "selected_files")
selected_files = driver.find_element(By.ID, f"selected_files{suffix}")
assert selected_files.text == exp_name

state = await upload_file.get_state(token)
if secondary:
# only the secondary form tracks progress and chain events
assert state.event_order.count("upload_progress") == 1
assert state.event_order.count("chain_event") == 1


@pytest.mark.asyncio
async def test_upload_file_multiple(tmp_path, upload_file: AppHarness, driver):
Expand Down Expand Up @@ -186,13 +259,17 @@ async def get_file_data():
assert file_data[exp_name] == exp_contents


def test_clear_files(tmp_path, upload_file: AppHarness, driver):
@pytest.mark.parametrize("secondary", [False, True])
def test_clear_files(
tmp_path, upload_file: AppHarness, driver: WebDriver, secondary: bool
):
"""Select then clear several file uploads and check that they are cleared.
Args:
tmp_path: pytest tmp_path fixture
upload_file: harness for UploadFile app.
driver: WebDriver instance.
secondary: whether to use the secondary upload form.
"""
assert upload_file.app_instance is not None
token_input = driver.find_element(By.ID, "token")
Expand All @@ -201,9 +278,13 @@ def test_clear_files(tmp_path, upload_file: AppHarness, driver):
token = upload_file.poll_for_value(token_input)
assert token is not None

upload_box = driver.find_element(By.XPATH, "//input[@type='file']")
suffix = "_secondary" if secondary else ""

upload_box = driver.find_elements(By.XPATH, "//input[@type='file']")[
1 if secondary else 0
]
assert upload_box
upload_button = driver.find_element(By.ID, "upload_button")
upload_button = driver.find_element(By.ID, f"upload_button{suffix}")
assert upload_button

exp_files = {
Expand All @@ -219,13 +300,56 @@ def test_clear_files(tmp_path, upload_file: AppHarness, driver):
time.sleep(0.2)

# check that the selected files are displayed
selected_files = driver.find_element(By.ID, "selected_files")
selected_files = driver.find_element(By.ID, f"selected_files{suffix}")
assert selected_files.text == "\n".join(exp_files)

clear_button = driver.find_element(By.ID, "clear_button")
clear_button = driver.find_element(By.ID, f"clear_button{suffix}")
assert clear_button
clear_button.click()

# check that the selected files are cleared
selected_files = driver.find_element(By.ID, "selected_files")
selected_files = driver.find_element(By.ID, f"selected_files{suffix}")
assert selected_files.text == ""


# TODO: drag and drop directory
# https://gist.github.com/florentbr/349b1ab024ca9f3de56e6bf8af2ac69e


@pytest.mark.asyncio
async def test_cancel_upload(tmp_path, upload_file: AppHarness, driver: WebDriver):
"""Submit a large file upload and cancel it.
Args:
tmp_path: pytest tmp_path fixture
upload_file: harness for UploadFile app.
driver: WebDriver instance.
"""
assert upload_file.app_instance is not None
token_input = driver.find_element(By.ID, "token")
assert token_input
# wait for the backend connection to send the token
token = upload_file.poll_for_value(token_input)
assert token is not None

upload_box = driver.find_elements(By.XPATH, "//input[@type='file']")[1]
upload_button = driver.find_element(By.ID, f"upload_button_secondary")
cancel_button = driver.find_element(By.ID, f"cancel_button_secondary")

exp_name = "large.txt"
target_file = tmp_path / exp_name
with target_file.open("wb") as f:
f.seek(1024 * 1024 * 256)
f.write(b"0")

upload_box.send_keys(str(target_file))
upload_button.click()
await asyncio.sleep(0.3)
cancel_button.click()

# look up the backend state and assert on progress
state = await upload_file.get_state(token)
assert state.progress_dicts
assert exp_name not in state._file_data

target_file.unlink()
Loading

0 comments on commit 4721c4f

Please sign in to comment.