|
1 | 1 | """Tests creating a custom extension"""
|
2 | 2 |
|
3 |
| -import unittest |
| 3 | +from collections.abc import Generator |
4 | 4 | from datetime import datetime
|
5 | 5 | from typing import Any, Generic, TypeVar, cast
|
6 | 6 |
|
@@ -119,81 +119,85 @@ def migrate(
|
119 | 119 | super().migrate(obj, version, info)
|
120 | 120 |
|
121 | 121 |
|
122 |
| -class CustomExtensionTest(unittest.TestCase): |
123 |
| - def setUp(self) -> None: |
124 |
| - pystac.EXTENSION_HOOKS.add_extension_hooks(CustomExtensionHooks()) |
125 |
| - |
126 |
| - def tearDown(self) -> None: |
127 |
| - pystac.EXTENSION_HOOKS.remove_extension_hooks(SCHEMA_URI) |
128 |
| - |
129 |
| - def test_add_to_item_asset(self) -> None: |
130 |
| - item = Item("an-id", None, None, datetime.now(), {}) |
131 |
| - item.add_asset("foo", Asset("http://pystac.test/asset.tif")) |
132 |
| - custom = CustomExtension.ext(item.assets["foo"], add_if_missing=True) |
133 |
| - assert CustomExtension.has_extension(item) |
134 |
| - custom.apply("bar") |
135 |
| - item_as_dict = item.to_dict() |
136 |
| - assert item_as_dict["assets"]["foo"]["test:prop"] == "bar" |
137 |
| - |
138 |
| - def test_add_to_item(self) -> None: |
139 |
| - item = Item("an-id", None, None, datetime.now(), {}) |
140 |
| - custom = CustomExtension.ext(item, add_if_missing=True) |
141 |
| - assert CustomExtension.has_extension(item) |
142 |
| - custom.test_prop = "foo" |
143 |
| - item_as_dict = item.to_dict() |
144 |
| - assert item_as_dict["properties"]["test:prop"] == "foo" |
145 |
| - |
146 |
| - def test_add_to_catalog(self) -> None: |
147 |
| - catalog = Catalog("an-id", "a description") |
148 |
| - custom = CustomExtension.ext(catalog, add_if_missing=True) |
149 |
| - assert CustomExtension.has_extension(catalog) |
150 |
| - custom.test_prop = "foo" |
151 |
| - catalog_as_dict = catalog.to_dict() |
152 |
| - assert catalog_as_dict["test:prop"] == "foo" |
153 |
| - |
154 |
| - def test_add_to_collection(self) -> None: |
155 |
| - collection = Collection( |
156 |
| - "an-id", |
157 |
| - "a description", |
158 |
| - extent=Extent( |
159 |
| - spatial=SpatialExtent([-180.0, -90.0, 180.0, 90.0]), |
160 |
| - temporal=TemporalExtent([[datetime.now(), None]]), |
161 |
| - ), |
162 |
| - ) |
163 |
| - custom = CustomExtension.ext(collection, add_if_missing=True) |
164 |
| - assert CustomExtension.has_extension(collection) |
165 |
| - custom.test_prop = "foo" |
166 |
| - collection_as_dict = collection.to_dict() |
167 |
| - assert collection_as_dict["test:prop"] == "foo" |
168 |
| - |
169 |
| - def test_add_to_collection_asset(self) -> None: |
170 |
| - collection = Collection( |
171 |
| - "an-id", |
172 |
| - "a description", |
173 |
| - extent=Extent( |
174 |
| - spatial=SpatialExtent([-180.0, -90.0, 180.0, 90.0]), |
175 |
| - temporal=TemporalExtent([[datetime.now(), None]]), |
176 |
| - ), |
177 |
| - ) |
178 |
| - collection.add_asset("foo", Asset("http://pystac.test/asset.tif")) |
179 |
| - custom = CustomExtension.ext(collection.assets["foo"], add_if_missing=True) |
180 |
| - assert CustomExtension.has_extension(collection) |
181 |
| - custom.test_prop = "bar" |
182 |
| - collection_as_dict = collection.to_dict() |
183 |
| - assert collection_as_dict["assets"]["foo"]["test:prop"] == "bar" |
184 |
| - |
185 |
| - def test_ext_non_stac_object(self) -> None: |
186 |
| - with pytest.raises(ExtensionTypeError): |
187 |
| - CustomExtension.ext({}) # type: ignore |
188 |
| - |
189 |
| - def test_migrates(self) -> None: |
190 |
| - item = Item("an-id", None, None, datetime.now(), {}) |
191 |
| - item_as_dict = item.to_dict() |
192 |
| - item_as_dict["stac_version"] = "1.0.0-rc.1" |
193 |
| - item_as_dict["stac_extensions"] = [ |
194 |
| - "https://example.com/v1.0/custom-schema.json" |
195 |
| - ] |
196 |
| - item_as_dict["properties"]["test:old-prop-name"] = "foo" |
197 |
| - item = Item.from_dict(item_as_dict, migrate=True) |
198 |
| - custom = CustomExtension.ext(item) |
199 |
| - assert custom.test_prop == "foo" |
| 122 | +@pytest.fixture |
| 123 | +def add_extension_hooks() -> Generator[None]: |
| 124 | + pystac.EXTENSION_HOOKS.add_extension_hooks(CustomExtensionHooks()) |
| 125 | + yield |
| 126 | + pystac.EXTENSION_HOOKS.remove_extension_hooks(SCHEMA_URI) |
| 127 | + |
| 128 | + |
| 129 | +def test_add_to_item_asset(add_extension_hooks: None) -> None: |
| 130 | + item = Item("an-id", None, None, datetime.now(), {}) |
| 131 | + item.add_asset("foo", Asset("http://pystac.test/asset.tif")) |
| 132 | + custom = CustomExtension.ext(item.assets["foo"], add_if_missing=True) |
| 133 | + assert CustomExtension.has_extension(item) |
| 134 | + custom.apply("bar") |
| 135 | + item_as_dict = item.to_dict() |
| 136 | + assert item_as_dict["assets"]["foo"]["test:prop"] == "bar" |
| 137 | + |
| 138 | + |
| 139 | +def test_add_to_item(add_extension_hooks: None) -> None: |
| 140 | + item = Item("an-id", None, None, datetime.now(), {}) |
| 141 | + custom = CustomExtension.ext(item, add_if_missing=True) |
| 142 | + assert CustomExtension.has_extension(item) |
| 143 | + custom.test_prop = "foo" |
| 144 | + item_as_dict = item.to_dict() |
| 145 | + assert item_as_dict["properties"]["test:prop"] == "foo" |
| 146 | + |
| 147 | + |
| 148 | +def test_add_to_catalog(add_extension_hooks: None) -> None: |
| 149 | + catalog = Catalog("an-id", "a description") |
| 150 | + custom = CustomExtension.ext(catalog, add_if_missing=True) |
| 151 | + assert CustomExtension.has_extension(catalog) |
| 152 | + custom.test_prop = "foo" |
| 153 | + catalog_as_dict = catalog.to_dict() |
| 154 | + assert catalog_as_dict["test:prop"] == "foo" |
| 155 | + |
| 156 | + |
| 157 | +def test_add_to_collection(add_extension_hooks: None) -> None: |
| 158 | + collection = Collection( |
| 159 | + "an-id", |
| 160 | + "a description", |
| 161 | + extent=Extent( |
| 162 | + spatial=SpatialExtent([-180.0, -90.0, 180.0, 90.0]), |
| 163 | + temporal=TemporalExtent([[datetime.now(), None]]), |
| 164 | + ), |
| 165 | + ) |
| 166 | + custom = CustomExtension.ext(collection, add_if_missing=True) |
| 167 | + assert CustomExtension.has_extension(collection) |
| 168 | + custom.test_prop = "foo" |
| 169 | + collection_as_dict = collection.to_dict() |
| 170 | + assert collection_as_dict["test:prop"] == "foo" |
| 171 | + |
| 172 | + |
| 173 | +def test_add_to_collection_asset(add_extension_hooks: None) -> None: |
| 174 | + collection = Collection( |
| 175 | + "an-id", |
| 176 | + "a description", |
| 177 | + extent=Extent( |
| 178 | + spatial=SpatialExtent([-180.0, -90.0, 180.0, 90.0]), |
| 179 | + temporal=TemporalExtent([[datetime.now(), None]]), |
| 180 | + ), |
| 181 | + ) |
| 182 | + collection.add_asset("foo", Asset("http://pystac.test/asset.tif")) |
| 183 | + custom = CustomExtension.ext(collection.assets["foo"], add_if_missing=True) |
| 184 | + assert CustomExtension.has_extension(collection) |
| 185 | + custom.test_prop = "bar" |
| 186 | + collection_as_dict = collection.to_dict() |
| 187 | + assert collection_as_dict["assets"]["foo"]["test:prop"] == "bar" |
| 188 | + |
| 189 | + |
| 190 | +def test_ext_non_stac_object(add_extension_hooks: None) -> None: |
| 191 | + with pytest.raises(ExtensionTypeError): |
| 192 | + CustomExtension.ext({}) # type: ignore |
| 193 | + |
| 194 | + |
| 195 | +def test_migrates(add_extension_hooks: None) -> None: |
| 196 | + item = Item("an-id", None, None, datetime.now(), {}) |
| 197 | + item_as_dict = item.to_dict() |
| 198 | + item_as_dict["stac_version"] = "1.0.0-rc.1" |
| 199 | + item_as_dict["stac_extensions"] = ["https://example.com/v1.0/custom-schema.json"] |
| 200 | + item_as_dict["properties"]["test:old-prop-name"] = "foo" |
| 201 | + item = Item.from_dict(item_as_dict, migrate=True) |
| 202 | + custom = CustomExtension.ext(item) |
| 203 | + assert custom.test_prop == "foo" |
0 commit comments