diff --git a/app/packages/components/src/components/PillBadge/PillBadge.tsx b/app/packages/components/src/components/PillBadge/PillBadge.tsx
new file mode 100644
index 0000000000..362bc80b7a
--- /dev/null
+++ b/app/packages/components/src/components/PillBadge/PillBadge.tsx
@@ -0,0 +1,149 @@
+import React, { useState } from "react";
+import CircleIcon from "@mui/icons-material/Circle";
+import { Chip, FormControl, MenuItem, Select } from "@mui/material";
+
+const PillBadge = ({
+ text,
+ color = "default",
+ variant = "filled",
+ showIcon = true,
+}: {
+ text: string | string[] | [string, string][];
+ color?: string;
+ variant?: "filled" | "outlined";
+ showIcon?: boolean;
+}) => {
+ const getInitialChipSelection = (
+ text: string | string[] | [string, string][]
+ ) => {
+ if (typeof text === "string") return text;
+ if (Array.isArray(text)) {
+ if (Array.isArray(text[0])) return text[0][0];
+ return text[0];
+ }
+ return "";
+ };
+
+ const getInitialChipColor = (
+ text: string | string[] | [string, string][],
+ color?: string
+ ) => {
+ if (typeof text === "string") return color;
+ if (Array.isArray(text)) {
+ if (Array.isArray(text[0])) return text[0][1];
+ return color || "default";
+ }
+ return "default";
+ };
+
+ const [chipSelection, setChipSelection] = useState(
+ getInitialChipSelection(text)
+ );
+ const [chipColor, setChipColor] = useState(getInitialChipColor(text, color));
+
+ const COLORS: { [key: string]: string } = {
+ default: "#999999",
+ primary: "#FFB682",
+ error: "error",
+ warning: "warning",
+ info: "info",
+ success: "#8BC18D",
+ };
+
+ const chipStyle: { [key: string]: string | number } = {
+ color: COLORS[chipColor || "default"] || COLORS.default,
+ fontWeight: 500,
+ paddingLeft: 1,
+ };
+
+ return (
+
+ {typeof text === "string" ? (
+
+ ) : undefined
+ }
+ label={text}
+ sx={{
+ ...chipStyle,
+ "& .MuiChip-icon": {
+ marginRight: "-7px",
+ },
+ "& .MuiChip-label": {
+ marginBottom: "1px",
+ },
+ }}
+ variant={variant as "filled" | "outlined" | undefined}
+ />
+ ) : (
+
+
+ ) : undefined
+ }
+ label={
+ Array.isArray(text) && text.length > 0 && Array.isArray(text[0]) ? (
+
+ ) : (
+
+ )
+ }
+ sx={{
+ ...chipStyle,
+ "& .MuiChip-icon": {
+ marginRight: "-7px",
+ },
+ "& .MuiChip-label": {
+ marginBottom: "1px",
+ },
+ "& .MuiInput-input:focus": {
+ backgroundColor: "inherit",
+ },
+ }}
+ variant={variant as "filled" | "outlined" | undefined}
+ >
+
+ )}
+
+ );
+};
+
+export default PillBadge;
diff --git a/app/packages/components/src/components/PillBadge/index.ts b/app/packages/components/src/components/PillBadge/index.ts
new file mode 100644
index 0000000000..45ed630ea2
--- /dev/null
+++ b/app/packages/components/src/components/PillBadge/index.ts
@@ -0,0 +1 @@
+export { default as PillBadge } from "./PillBadge";
diff --git a/app/packages/core/src/plugins/SchemaIO/components/PillBadgeView.tsx b/app/packages/core/src/plugins/SchemaIO/components/PillBadgeView.tsx
new file mode 100644
index 0000000000..787fabbd7d
--- /dev/null
+++ b/app/packages/core/src/plugins/SchemaIO/components/PillBadgeView.tsx
@@ -0,0 +1,22 @@
+import { Box } from "@mui/material";
+import React from "react";
+import { getComponentProps } from "../utils";
+import PillBadge from "@fiftyone/components/src/components/PillBadge/PillBadge";
+
+export default function PillBadgeView(props) {
+ const { schema } = props;
+ const { view = {} } = schema;
+ const { text, color, variant, showIcon } = view;
+
+ return (
+
+
+
+ );
+}
diff --git a/app/packages/core/src/plugins/SchemaIO/components/index.ts b/app/packages/core/src/plugins/SchemaIO/components/index.ts
index bb0fca6f6e..fefdc1c759 100644
--- a/app/packages/core/src/plugins/SchemaIO/components/index.ts
+++ b/app/packages/core/src/plugins/SchemaIO/components/index.ts
@@ -49,3 +49,4 @@ export { default as TextFieldView } from "./TextFieldView";
export { default as TupleView } from "./TupleView";
export { default as UnsupportedView } from "./UnsupportedView";
export { default as FrameLoaderView } from "./FrameLoaderView";
+export { default as PillBadgeView } from "./PillBadgeView";
diff --git a/fiftyone/factory/repo_factory.py b/fiftyone/factory/repo_factory.py
index 2b031588da..314163544e 100644
--- a/fiftyone/factory/repo_factory.py
+++ b/fiftyone/factory/repo_factory.py
@@ -13,6 +13,10 @@
DelegatedOperationRepo,
MongoDelegatedOperationRepo,
)
+from fiftyone.factory.repos.execution_store import (
+ ExecutionStoreRepo,
+ MongoExecutionStoreRepo,
+)
_db: Database = None
@@ -44,3 +48,18 @@ def delegated_operation_repo() -> DelegatedOperationRepo:
return RepositoryFactory.repos[
MongoDelegatedOperationRepo.COLLECTION_NAME
]
+
+ @staticmethod
+ def execution_store_repo() -> ExecutionStoreRepo:
+ """Factory method for execution store repository."""
+ if (
+ MongoExecutionStoreRepo.COLLECTION_NAME
+ not in RepositoryFactory.repos
+ ):
+ RepositoryFactory.repos[
+ MongoExecutionStoreRepo.COLLECTION_NAME
+ ] = MongoExecutionStoreRepo(
+ collection=_get_db()[MongoExecutionStoreRepo.COLLECTION_NAME]
+ )
+
+ return RepositoryFactory.repos[MongoExecutionStoreRepo.COLLECTION_NAME]
diff --git a/fiftyone/factory/repos/execution_store.py b/fiftyone/factory/repos/execution_store.py
new file mode 100644
index 0000000000..69fa78b077
--- /dev/null
+++ b/fiftyone/factory/repos/execution_store.py
@@ -0,0 +1,128 @@
+"""
+Execution store repository.
+"""
+
+import datetime
+from pymongo.collection import Collection
+from fiftyone.operators.store.models import StoreDocument, KeyDocument
+
+
+def _where(store_name, key=None):
+ query = {"store_name": store_name}
+ if key is not None:
+ query["key"] = key
+ return query
+
+
+class ExecutionStoreRepo:
+ """Base class for execution store repositories."""
+
+ COLLECTION_NAME = "execution_store"
+
+ def __init__(self, collection: Collection):
+ self._collection = collection
+
+ def create_store(self, store_name, permissions=None) -> StoreDocument:
+ """Creates a store in the execution store."""
+ store_doc = StoreDocument(
+ store_name=store_name, permissions=permissions
+ )
+ self._collection.insert_one(store_doc.dict())
+ return store_doc
+
+ def list_stores(self) -> list[str]:
+ """Lists all stores in the execution store."""
+ # ensure that only store_name is returned, and only unique values
+ return self._collection.distinct("store_name")
+
+ def set_key(self, store_name, key, value, ttl=None) -> KeyDocument:
+ """Sets or updates a key in the specified store"""
+ now = datetime.datetime.now()
+ expiration = KeyDocument.get_expiration(ttl)
+ key_doc = KeyDocument(
+ store_name=store_name, key=key, value=value, updated_at=now
+ )
+
+ # Prepare the update operations
+ update_fields = {
+ "$set": key_doc.dict(
+ exclude={"created_at", "expires_at", "store_name", "key"}
+ ),
+ "$setOnInsert": {
+ "store_name": store_name,
+ "key": key,
+ "created_at": now,
+ "expires_at": expiration if ttl else None,
+ },
+ }
+
+ # Perform the upsert operation
+ result = self._collection.update_one(
+ _where(store_name, key), update_fields, upsert=True
+ )
+
+ if result.upserted_id:
+ key_doc.created_at = now
+ else:
+ key_doc.updated_at = now
+
+ return key_doc
+
+ def get_key(self, store_name, key) -> KeyDocument:
+ """Gets a key from the specified store."""
+ raw_key_doc = self._collection.find_one(_where(store_name, key))
+ key_doc = KeyDocument(**raw_key_doc) if raw_key_doc else None
+ return key_doc
+
+ def list_keys(self, store_name) -> list[str]:
+ """Lists all keys in the specified store."""
+ keys = self._collection.find(_where(store_name), {"key": 1})
+ return [key["key"] for key in keys]
+
+ def update_ttl(self, store_name, key, ttl) -> bool:
+ """Updates the TTL for a key."""
+ expiration = KeyDocument.get_expiration(ttl)
+ result = self._collection.update_one(
+ _where(store_name, key), {"$set": {"expires_at": expiration}}
+ )
+ return result.modified_count > 0
+
+ def delete_key(self, store_name, key) -> bool:
+ """Deletes the document that matches the store name and key."""
+ result = self._collection.delete_one(_where(store_name, key))
+ return result.deleted_count > 0
+
+ def delete_store(self, store_name) -> int:
+ """Deletes the entire store."""
+ result = self._collection.delete_many(_where(store_name))
+ return result.deleted_count
+
+
+class MongoExecutionStoreRepo(ExecutionStoreRepo):
+ """MongoDB implementation of execution store repository."""
+
+ COLLECTION_NAME = "execution_store"
+
+ def __init__(self, collection: Collection):
+ super().__init__(collection)
+ self._create_indexes()
+
+ def _create_indexes(self):
+ indices = self._collection.list_indexes()
+ expires_at_name = "expires_at"
+ store_name_name = "store_name"
+ key_name = "key"
+ full_key_name = "store_name_and_key"
+ if expires_at_name not in indices:
+ self._collection.create_index(
+ expires_at_name, name=expires_at_name, expireAfterSeconds=0
+ )
+ if full_key_name not in indices:
+ self._collection.create_index(
+ [(store_name_name, 1), (key_name, 1)],
+ name=full_key_name,
+ unique=True,
+ )
+ for name in [store_name_name, key_name]:
+ if name not in indices:
+ self._collection.create_index(name, name=name)
diff --git a/fiftyone/operators/executor.py b/fiftyone/operators/executor.py
index e0e6ac784c..0f9ab5a9f8 100644
--- a/fiftyone/operators/executor.py
+++ b/fiftyone/operators/executor.py
@@ -852,6 +852,20 @@ def set_progress(self, progress=None, label=None):
else:
self.log(f"Progress: {progress} - {label}")
+ # TODO resolve circular import so this can have a type
+ def create_store(self, store_name):
+ """Creates a new store with the specified name.
+
+ Args:
+ store_name: the name of the store
+
+ Returns:
+ a :class:`fiftyone.operators.store.ExecutionStore`
+ """
+ from fiftyone.operators.store import ExecutionStore
+
+ return ExecutionStore.create(store_name)
+
def serialize(self):
"""Serializes the execution context.
diff --git a/fiftyone/operators/store/__init__.py b/fiftyone/operators/store/__init__.py
new file mode 100644
index 0000000000..50f9f3b400
--- /dev/null
+++ b/fiftyone/operators/store/__init__.py
@@ -0,0 +1,18 @@
+"""
+FiftyOne execution store module.
+
+| Copyright 2017-2024, Voxel51, Inc.
+| `voxel51.com `_
+|
+"""
+
+from .service import ExecutionStoreService
+from .store import ExecutionStore
+from .models import StoreDocument, KeyDocument
+
+__all__ = [
+ "ExecutionStoreService",
+ "StoreDocument",
+ "KeyDocument",
+ "ExecutionStore",
+]
diff --git a/fiftyone/operators/store/models.py b/fiftyone/operators/store/models.py
new file mode 100644
index 0000000000..5fd6c89e21
--- /dev/null
+++ b/fiftyone/operators/store/models.py
@@ -0,0 +1,35 @@
+"""
+Store and key models for the execution store.
+"""
+
+from pydantic import BaseModel, Field
+from typing import Optional, Dict, Any
+import datetime
+
+
+class KeyDocument(BaseModel):
+ """Model representing a key in the store."""
+
+ store_name: str
+ key: str
+ value: Any
+ created_at: datetime.datetime = Field(
+ default_factory=datetime.datetime.now
+ )
+ updated_at: Optional[datetime.datetime] = None
+ expires_at: Optional[datetime.datetime] = None
+
+ @staticmethod
+ def get_expiration(ttl: Optional[int]) -> Optional[datetime.datetime]:
+ """Gets the expiration date for a key with the given TTL."""
+ if ttl is None:
+ return None
+
+ return datetime.datetime.now() + datetime.timedelta(seconds=ttl)
+
+
+class StoreDocument(KeyDocument):
+ """Model representing a Store."""
+
+ key: str = "__store__"
+ value: Optional[Dict[str, Any]] = None
diff --git a/fiftyone/operators/store/service.py b/fiftyone/operators/store/service.py
new file mode 100644
index 0000000000..dc90039985
--- /dev/null
+++ b/fiftyone/operators/store/service.py
@@ -0,0 +1,115 @@
+"""
+FiftyOne execution store service.
+
+| Copyright 2017-2024, Voxel51, Inc.
+| `voxel51.com `_
+|
+"""
+
+import logging
+from fiftyone.factory.repo_factory import RepositoryFactory
+
+logger = logging.getLogger(__name__)
+
+
+class ExecutionStoreService(object):
+ """Service for managing execution store operations."""
+
+ def __init__(self, repo=None):
+ if repo is None:
+ repo = RepositoryFactory.execution_store_repo()
+
+ self._repo = repo
+
+ def create_store(self, store_name):
+ """Creates a new store with the specified name.
+
+ Args:
+ store_name: the name of the store
+
+ Returns:
+ a :class:`fiftyone.store.models.StoreDocument`
+ """
+ return self._repo.create_store(store_name=store_name)
+
+ def set_key(self, store_name, key, value, ttl=None):
+ """Sets the value of a key in the specified store.
+
+ Args:
+ store_name: the name of the store
+ key: the key to set
+ value: the value to set
+ ttl (None): an optional TTL in seconds
+
+ Returns:
+ a :class:`fiftyone.store.models.KeyDocument`
+ """
+ return self._repo.set_key(
+ store_name=store_name, key=key, value=value, ttl=ttl
+ )
+
+ def get_key(self, store_name, key):
+ """Retrieves the value of a key from the specified store.
+
+ Args:
+ store_name: the name of the store
+ key: the key to retrieve
+
+ Returns:
+ a :class:`fiftyone.store.models.KeyDocument`
+ """
+ return self._repo.get_key(store_name=store_name, key=key)
+
+ def delete_key(self, store_name, key):
+ """Deletes the specified key from the store.
+
+ Args:
+ store_name: the name of the store
+ key: the key to delete
+
+ Returns:
+ `True` if the key was deleted, `False` otherwise
+ """
+ return self._repo.delete_key(store_name=store_name, key=key)
+
+ def update_ttl(self, store_name, key, new_ttl):
+ """Updates the TTL of the specified key in the store.
+
+ Args:
+ store_name: the name of the store
+ key: the key to update the TTL for
+ new_ttl: the new TTL in milliseconds
+
+ Returns:
+ a :class:`fiftyone.store.models.KeyDocument`
+ """
+ return self._repo.update_ttl(
+ store_name=store_name, key=key, ttl=new_ttl
+ )
+
+ def list_stores(self):
+ """Lists all stores matching the given criteria.
+
+ Returns:
+ a list of :class:`fiftyone.store.models.StoreDocument`
+ """
+ return self._repo.list_stores()
+
+ def delete_store(self, store_name):
+ """Deletes the specified store.
+
+ Args:
+ store_name: the name of the store
+
+ Returns:
+ a :class:`fiftyone.store.models.StoreDocument`
+ """
+ return self._repo.delete_store(store_name=store_name)
+
+ def list_keys(self, store_name):
+ """Lists all keys in the specified store.
+
+ Args:
+ store_name: the name of the store
+ """
+ return self._repo.list_keys(store_name)
diff --git a/fiftyone/operators/store/store.py b/fiftyone/operators/store/store.py
new file mode 100644
index 0000000000..15e0ab3a01
--- /dev/null
+++ b/fiftyone/operators/store/store.py
@@ -0,0 +1,115 @@
+"""
+FiftyOne execution store.
+
+| Copyright 2017-2024, Voxel51, Inc.
+| `voxel51.com `_
+|
+"""
+
+import logging
+from fiftyone.factory.repo_factory import RepositoryFactory
+from fiftyone.operators.store.service import ExecutionStoreService
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class ExecutionStore:
+ @staticmethod
+ def create(store_name: str) -> "ExecutionStore":
+ return ExecutionStore(store_name, ExecutionStoreService())
+
+ def __init__(self, store_name: str, store_service: ExecutionStoreService):
+ """
+ Args:
+ store_name (str): The name of the store.
+ store_service (ExecutionStoreService): The store service instance.
+ """
+ self.store_name: str = store_name
+ self._store_service: ExecutionStoreService = store_service
+
+ def list_all_stores(self) -> list[str]:
+ """Lists all stores in the execution store.
+
+ Returns:
+ list: A list of store names.
+ """
+ return self._store_service.list_stores()
+
+ def get(self, key: str) -> Optional[Any]:
+ """Retrieves a value from the store by its key.
+
+ Args:
+ key (str): The key to retrieve the value for.
+
+ Returns:
+ Optional[Any]: The value stored under the given key, or None if not found.
+ """
+ key_doc = self._store_service.get_key(self.store_name, key)
+ if key_doc is None:
+ return None
+ return key_doc.value
+
+ def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
+ """Sets a value in the store with an optional TTL.
+
+ Args:
+ key (str): The key to store the value under.
+ value (Any): The value to store.
+ ttl (Optional[int], optional): The time-to-live in seconds. Defaults to None.
+ """
+ self._store_service.set_key(self.store_name, key, value, ttl)
+
+ def delete(self, key: str) -> None:
+ """Deletes a key from the store.
+
+ Args:
+ key (str): The key to delete.
+
+ Returns:
+ bool: True if the key was deleted, False otherwise.
+ """
+ return self._store_service.delete_key(self.store_name, key)
+
+ def has(self, key: str) -> bool:
+ """Checks if the store has a specific key.
+
+ Args:
+ key (str): The key to check.
+
+ Returns:
+ bool: True if the key exists, False otherwise.
+ """
+ return self._store_service.has_key(self.store_name, key)
+
+ def clear(self) -> None:
+ """Clears all the data in the store."""
+ self._store_service.delete_store(self.store_name)
+
+ def update_ttl(self, key: str, new_ttl: int) -> None:
+ """Updates the TTL for a specific key.
+
+ Args:
+ key (str): The key to update the TTL for.
+ new_ttl (int): The new TTL in milliseconds.
+ """
+ self._store_service.update_ttl(self.store_name, key, new_ttl)
+
+ def get_ttl(self, key: str) -> Optional[int]:
+ """Retrieves the TTL for a specific key.
+
+ Args:
+ key (str): The key to get the TTL for.
+
+ Returns:
+ Optional[int]: The TTL in milliseconds, or None if the key does not have a TTL.
+ """
+ return self._store_service.get_ttl(self.store_name, key)
+
+ def list_keys(self) -> list[str]:
+ """Lists all keys in the store.
+
+ Returns:
+ list: A list of keys in the store.
+ """
+ return self._store_service.list_keys(self.store_name)
diff --git a/fiftyone/operators/types.py b/fiftyone/operators/types.py
index a22bd01230..5117ca0c41 100644
--- a/fiftyone/operators/types.py
+++ b/fiftyone/operators/types.py
@@ -1475,6 +1475,20 @@ def __init__(self, **kwargs):
super().__init__(**kwargs)
+class PillBadgeView(ReadOnlyView):
+ """Displays a pill shaped badge.
+
+ Args:
+ text ("Reviewed" | ["Reviewed", "Not Reviewed"] | [["Not Started", "primary"], ["Reviewed", "success"], ["In Review", "warning"]): a label or set of label options with or without a color for the pill badge
+ color ("primary"): the color of the pill
+ variant ("outlined"): the variant of the pill
+ show_icon (False | True): whether to display indicator icon
+ """
+
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+
class PlotlyView(View):
"""Displays a Plotly chart.
diff --git a/tests/unittests/execution_store_tests.py b/tests/unittests/execution_store_tests.py
new file mode 100644
index 0000000000..470d399026
--- /dev/null
+++ b/tests/unittests/execution_store_tests.py
@@ -0,0 +1,228 @@
+"""
+FiftyOne execution store related unit tests.
+
+| Copyright 2017-2024, Voxel51, Inc.
+| `voxel51.com `_
+|
+"""
+
+import time
+import unittest
+from unittest import mock
+from unittest.mock import patch, MagicMock, ANY, Mock
+import datetime
+
+import fiftyone.operators as foo
+from fiftyone.operators.store import ExecutionStoreService
+from fiftyone.operators.store.models import StoreDocument, KeyDocument
+from fiftyone.factory.repo_factory import ExecutionStoreRepo
+from fiftyone.operators.store import ExecutionStore
+from fiftyone.operators.store import ExecutionStoreService
+
+EPSILON = 0.1
+
+
+class IsDateTime:
+ def __eq__(self, other):
+ return isinstance(other, datetime.datetime)
+
+
+def assert_delta_seconds_approx(time_delta, seconds, epsilon=EPSILON):
+ assert abs(time_delta.total_seconds() - seconds) < epsilon
+
+
+class TeskKeyDocument(unittest.TestCase):
+ def test_get_expiration(self):
+ ttl = 1
+ now = datetime.datetime.now()
+ expiration = KeyDocument.get_expiration(ttl)
+ time_delta = expiration - now
+ assert_delta_seconds_approx(time_delta, ttl)
+ assert isinstance(expiration, datetime.datetime)
+
+ def test_get_expiration_none(self):
+ ttl = None
+ expiration = KeyDocument.get_expiration(ttl)
+ assert expiration is None
+
+
+class ExecutionStoreServiceIntegrationTests(unittest.TestCase):
+ def setUp(self) -> None:
+ self.mock_collection = MagicMock()
+ self.store_repo = ExecutionStoreRepo(self.mock_collection)
+ self.store_service = ExecutionStoreService(self.store_repo)
+
+ def test_set_key(self):
+ self.store_repo.set_key(
+ "widgets",
+ "widget_1",
+ {"name": "Widget One", "value": 100},
+ ttl=60000,
+ )
+ self.mock_collection.update_one.assert_called_once()
+ self.mock_collection.update_one.assert_called_with(
+ {"store_name": "widgets", "key": "widget_1"},
+ {
+ "$set": {
+ "value": {"name": "Widget One", "value": 100},
+ "updated_at": IsDateTime(),
+ },
+ "$setOnInsert": {
+ "store_name": "widgets",
+ "key": "widget_1",
+ "created_at": IsDateTime(),
+ "expires_at": IsDateTime(),
+ },
+ },
+ upsert=True,
+ )
+
+ def test_get_key(self):
+ self.mock_collection.find_one.return_value = {
+ "store_name": "widgets",
+ "key": "widget_1",
+ "value": {"name": "Widget One", "value": 100},
+ "created_at": time.time(),
+ "updated_at": time.time(),
+ "expires_at": time.time() + 60000,
+ }
+ self.store_service.get_key(store_name="widgets", key="widget_1")
+ self.mock_collection.find_one.assert_called_once()
+ self.mock_collection.find_one.assert_called_with(
+ {"store_name": "widgets", "key": "widget_1"}
+ )
+
+ def test_create_store(self):
+ self.store_repo.create_store("widgets")
+ self.mock_collection.insert_one.assert_called_once()
+ self.mock_collection.insert_one.assert_called_with(
+ {
+ "store_name": "widgets",
+ "key": "__store__",
+ "value": None,
+ "created_at": IsDateTime(),
+ "updated_at": None,
+ "expires_at": None,
+ }
+ )
+
+ def test_delete_key(self):
+ self.mock_collection.delete_one.return_value = Mock(deleted_count=1)
+ self.store_repo.delete_key("widgets", "widget_1")
+ self.mock_collection.delete_one.assert_called_once()
+ self.mock_collection.delete_one.assert_called_with(
+ {"store_name": "widgets", "key": "widget_1"}
+ )
+
+ def test_update_ttl(self):
+ self.mock_collection.update_one.return_value = Mock(modified_count=1)
+ ttl_seconds = 60000
+ expected_expiration = KeyDocument.get_expiration(ttl_seconds)
+ self.store_repo.update_ttl("widgets", "widget_1", ttl_seconds)
+ self.mock_collection.update_one.assert_called_once()
+
+ actual_call = self.mock_collection.update_one.call_args
+ actual_expires_at = actual_call[0][1]["$set"]["expires_at"]
+ time_delta = expected_expiration - actual_expires_at
+ assert_delta_seconds_approx(time_delta, 0, epsilon=0.0001)
+
+ def test_delete_store(self):
+ self.mock_collection.delete_many.return_value = Mock(deleted_count=1)
+ self.store_repo.delete_store("widgets")
+ self.mock_collection.delete_many.assert_called_once()
+ self.mock_collection.delete_many.assert_called_with(
+ {"store_name": "widgets"}
+ )
+
+ def test_list_keys(self):
+ self.mock_collection.find.return_value = [
+ {"store_name": "widgets", "key": "widget_1"},
+ {"store_name": "widgets", "key": "widget_2"},
+ ]
+ keys = self.store_repo.list_keys("widgets")
+ assert keys == ["widget_1", "widget_2"]
+ self.mock_collection.find.assert_called_once()
+ self.mock_collection.find.assert_called_with(
+ {"store_name": "widgets"}, {"key": 1}
+ )
+
+ def test_list_stores(self):
+ self.mock_collection.distinct.return_value = ["widgets", "gadgets"]
+ stores = self.store_repo.list_stores()
+ assert stores == ["widgets", "gadgets"]
+ self.mock_collection.distinct.assert_called_once()
+ self.mock_collection.distinct.assert_called_with("store_name")
+
+
+class TestExecutionStoreIntegration(unittest.TestCase):
+ def setUp(self) -> None:
+ self.mock_collection = MagicMock()
+ self.store_repo = ExecutionStoreRepo(self.mock_collection)
+ self.store_service = ExecutionStoreService(self.store_repo)
+ self.store = ExecutionStore("mock_store", self.store_service)
+
+ def test_set(self):
+ self.store.set(
+ "widget_1", {"name": "Widget One", "value": 100}, ttl=60000
+ )
+ self.mock_collection.update_one.assert_called_once()
+ self.mock_collection.update_one.assert_called_with(
+ {"store_name": "mock_store", "key": "widget_1"},
+ {
+ "$set": {
+ "updated_at": IsDateTime(),
+ "value": {"name": "Widget One", "value": 100},
+ },
+ "$setOnInsert": {
+ "store_name": "mock_store",
+ "key": "widget_1",
+ "created_at": IsDateTime(),
+ "expires_at": IsDateTime(),
+ },
+ },
+ upsert=True,
+ )
+
+ def test_get(self):
+ self.mock_collection.find_one.return_value = {
+ "store_name": "mock_store",
+ "key": "widget_1",
+ "value": {"name": "Widget One", "value": 100},
+ "created_at": time.time(),
+ "updated_at": time.time(),
+ "expires_at": time.time() + 60000,
+ }
+ value = self.store.get("widget_1")
+ assert value == {"name": "Widget One", "value": 100}
+ self.mock_collection.find_one.assert_called_once()
+ self.mock_collection.find_one.assert_called_with(
+ {"store_name": "mock_store", "key": "widget_1"}
+ )
+
+ def test_list_keys(self):
+ self.mock_collection.find.return_value = [
+ {"store_name": "mock_store", "key": "widget_1"},
+ {"store_name": "mock_store", "key": "widget_2"},
+ ]
+ keys = self.store.list_keys()
+ assert keys == ["widget_1", "widget_2"]
+ self.mock_collection.find.assert_called_once()
+ self.mock_collection.find.assert_called_with(
+ {"store_name": "mock_store"}, {"key": 1}
+ )
+
+ def test_delete(self):
+ self.mock_collection.delete_one.return_value = Mock(deleted_count=1)
+ deleted = self.store.delete("widget_1")
+ assert deleted
+ self.mock_collection.delete_one.assert_called_once()
+ self.mock_collection.delete_one.assert_called_with(
+ {"store_name": "mock_store", "key": "widget_1"}
+ )
+
+ def test_clear(self):
+ self.store.clear()
+ self.mock_collection.delete_many.assert_called_once()
+ self.mock_collection.delete_many.assert_called_with(
+ {"store_name": "mock_store"}
+ )