120 lines
3.5 KiB
Python
120 lines
3.5 KiB
Python
"""Unit tests for bulk artifact creation functionality."""
|
|
|
|
import pytest
|
|
|
|
from skyvern.forge.sdk.artifact.manager import ArtifactBatchData, BulkArtifactCreationRequest
|
|
from skyvern.forge.sdk.artifact.models import ArtifactType
|
|
from skyvern.forge.sdk.db.models import ArtifactModel
|
|
|
|
|
|
def test_artifact_batch_data_with_data():
|
|
"""Test ArtifactBatchData with data field."""
|
|
model = ArtifactModel(
|
|
artifact_id="test-1",
|
|
artifact_type=ArtifactType.SCREENSHOT_LLM,
|
|
uri="s3://bucket/test",
|
|
organization_id="org-1",
|
|
)
|
|
|
|
batch_data = ArtifactBatchData(
|
|
artifact_model=model,
|
|
data=b"test data",
|
|
)
|
|
|
|
assert batch_data.artifact_model == model
|
|
assert batch_data.data == b"test data"
|
|
assert batch_data.path is None
|
|
|
|
|
|
def test_artifact_batch_data_with_path():
|
|
"""Test ArtifactBatchData with path field."""
|
|
model = ArtifactModel(
|
|
artifact_id="test-1",
|
|
artifact_type=ArtifactType.SCREENSHOT_LLM,
|
|
uri="s3://bucket/test",
|
|
organization_id="org-1",
|
|
)
|
|
|
|
batch_data = ArtifactBatchData(
|
|
artifact_model=model,
|
|
path="/tmp/test.png",
|
|
)
|
|
|
|
assert batch_data.artifact_model == model
|
|
assert batch_data.data is None
|
|
assert batch_data.path == "/tmp/test.png"
|
|
|
|
|
|
def test_artifact_batch_data_with_both_raises_error():
|
|
"""Test that ArtifactBatchData raises error when both data and path are provided."""
|
|
model = ArtifactModel(
|
|
artifact_id="test-1",
|
|
artifact_type=ArtifactType.SCREENSHOT_LLM,
|
|
uri="s3://bucket/test",
|
|
organization_id="org-1",
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="Cannot specify both data and path"):
|
|
ArtifactBatchData(
|
|
artifact_model=model,
|
|
data=b"test data",
|
|
path="/tmp/test.png",
|
|
)
|
|
|
|
|
|
def test_bulk_artifact_creation_request():
|
|
"""Test BulkArtifactCreationRequest structure."""
|
|
model1 = ArtifactModel(
|
|
artifact_id="test-1",
|
|
artifact_type=ArtifactType.LLM_PROMPT,
|
|
uri="s3://bucket/test1",
|
|
organization_id="org-1",
|
|
)
|
|
model2 = ArtifactModel(
|
|
artifact_id="test-2",
|
|
artifact_type=ArtifactType.SCREENSHOT_LLM,
|
|
uri="s3://bucket/test2",
|
|
organization_id="org-1",
|
|
)
|
|
|
|
request = BulkArtifactCreationRequest(
|
|
artifacts=[
|
|
ArtifactBatchData(artifact_model=model1, data=b"data1"),
|
|
ArtifactBatchData(artifact_model=model2, data=b"data2"),
|
|
],
|
|
primary_key="task-123",
|
|
)
|
|
|
|
assert len(request.artifacts) == 2
|
|
assert request.primary_key == "task-123"
|
|
assert request.artifacts[0].artifact_model.artifact_id == "test-1"
|
|
assert request.artifacts[1].artifact_model.artifact_id == "test-2"
|
|
|
|
|
|
def test_bulk_artifact_creation_performance_benefit():
|
|
"""
|
|
Test to verify that bulk creation reduces database calls.
|
|
This is a conceptual test to document the performance improvement.
|
|
"""
|
|
# Before optimization: Creating N artifacts = N database INSERT calls
|
|
# After optimization: Creating N artifacts = 1 bulk INSERT call
|
|
|
|
num_artifacts = 10
|
|
|
|
# Simulate old approach (N individual inserts)
|
|
individual_insert_count = num_artifacts
|
|
|
|
# Simulate new approach (1 bulk insert)
|
|
bulk_insert_count = 1
|
|
|
|
# Assert that bulk insert is more efficient
|
|
assert bulk_insert_count < individual_insert_count
|
|
|
|
# The reduction ratio
|
|
reduction_ratio = individual_insert_count / bulk_insert_count
|
|
assert reduction_ratio == num_artifacts
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|