add test

Author: dikwickley

src/zarr/core/group.py

@@ -7,7 +7,7 @@
 import warnings
 from collections import defaultdict
 from dataclasses import asdict, dataclass, field, fields, replace
-from typing import TYPE_CHECKING, Literal, TypeVar, assert_never, cast, overload, Self
+from typing import TYPE_CHECKING, Literal, Self, TypeVar, assert_never, cast, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -55,6 +55,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, Generator, Iterable, Iterator
+    from types import TracebackType
     from typing import Any
 
     from zarr.core.array_spec import ArrayConfig, ArrayConfigLike
@@ -1755,7 +1756,9 @@ def open(
     def __enter__(self) -> Self:
         return self
 
-    def __exit__(self) -> None:  # noqa: PYI036
+    def __exit__(
+        self, typ: type[BaseException] | None, exc: BaseException | None, tb: TracebackType | None
+    ) -> None:
         self.store.close()
 
     def __getitem__(self, path: str) -> Array | Group:

tests/test_group.py

@@ -428,6 +428,28 @@ def test_group_len(store: Store, zarr_format: ZarrFormat) -> None:
     assert len(group) == 0
 
 
+def test_group_with_context_manager(store: Store, zarr_format: ZarrFormat, overwrite: bool) -> None:
+    spath = StorePath(store)
+
+    # attempt to open a group that does not exist.
+    with pytest.raises(FileNotFoundError):
+        with Group.open(store) as store:
+            pass
+
+    attrs = {"path": "foo"}
+
+    with Group.from_store(
+        store, attributes=attrs, zarr_format=zarr_format, overwrite=overwrite
+    ) as group:
+        assert store._is_open
+        assert group.attrs == attrs
+        assert group.metadata.zarr_format == zarr_format
+        assert group.store_path == spath
+
+    # Check if store was closed after exit.
+    assert not store._is_open
+
+
 def test_group_setitem(store: Store, zarr_format: ZarrFormat) -> None:
     """
     Test the `Group.__setitem__` method.