Coverage for tests/test_nibrs_store.py: 100%
197 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-30 18:01 +0000
1"""Tests of converting master files, loading tables, and building tables.
3Contents:
4 test_build_incidents: checks the built incidents.
5 test_build_other_tables: checks the other built tables.
6 test_build_victims_joins_offenders: checks the built victims.
7 test_convert_and_load: checks converting and loading.
8 test_convert_skips_converted_tables: checks converting again.
9 test_decode_later: checks decoding loaded codes.
10 test_download_checks_years: checks the years of downloads.
11 test_download_gives_up_without_progress: checks giving up downloads.
12 test_download_resumes_with_new_links: checks resuming downloads.
13 test_download_saves_files: checks saving downloads.
14 test_download_starts_over_if_range_is_ignored: checks restarting.
15 test_load_converts_missing_years: checks loading unconverted years.
16 test_load_keeps_states_and_columns: checks loading some states.
17 test_load_tables: checks loading several tables.
18 test_unknown_codes_are_kept: checks unknown codes.
20"""
22from __future__ import annotations
24import pathlib
25from typing import Any
27import amos
28import pandas as pd
29import pytest
31from crimepy import nibrs
32from crimepy.nibrs import cde, munge, tables
34""" Tests """
37def test_convert_and_load(data_folder: pathlib.Path) -> None:
38 """Checks that converted years load as a dataset with its label and groups.
40 Args:
41 data_folder: a data folder with the master files of 2004 and 2025.
43 """
44 converted = nibrs.convert([2004, 2025], data_folder, progress = False)
45 assert all(
46 p.is_file() for paths in converted.values() for p in paths.values())
47 loaded = nibrs.load(
48 'victims', '2004, 2025', folder = data_folder, label = 'victim_age',
49 task = 'regress', groups = 'ori')
50 assert isinstance(loaded, amos.Dataset)
51 assert loaded.label == 'victim_age'
52 assert loaded.task == 'regress'
53 assert loaded.groups == ['ori']
54 victims = loaded.data
55 assert sorted(victims['data_year'].unique()) == [2004, 2025]
56 assert len(victims) == 3
57 assert isinstance(victims['victim_sex'].dtype, pd.CategoricalDtype)
58 assert victims['victim_age'].dtype == 'Int64'
61def test_load_keeps_states_and_columns(data_folder: pathlib.Path) -> None:
62 """Checks that loading keeps only the states and columns asked for.
64 Args:
65 data_folder: a data folder with the master files of 2004 and 2025.
67 """
68 nibrs.convert('2004-2004', data_folder, progress = False)
69 victims = nibrs.load(
70 'victims', 2004, folder = data_folder, states = 'AZ',
71 columns = ['state_code', 'victim_age_code'], decode = False).data
72 assert list(victims.columns) == [
73 'state_code', 'state', 'victim_age_code', 'victim_age']
74 assert set(victims['state_code']) == {'02'}
75 empty = nibrs.load_frame(
76 'victims', 2004, folder = data_folder, states = 'AK')
77 assert empty.empty
78 with pytest.raises(KeyError):
79 nibrs.load('victims', 2004, folder = data_folder, columns = ['x'])
82def test_load_converts_missing_years(data_folder: pathlib.Path) -> None:
83 """Checks that loading converts years that were not converted.
85 Args:
86 data_folder: a data folder with the master files of 2004 and 2025.
88 """
89 offenders = nibrs.load(
90 'offenders', 2025, folder = data_folder, progress = False).data
91 race = offenders.iloc[0]['offender_race']
92 assert race == 'American Indian or Alaska Native'
93 with pytest.raises(FileNotFoundError):
94 nibrs.load('offenders', 2024, folder = data_folder, fetch = False)
97def test_convert_skips_converted_tables(data_folder: pathlib.Path) -> None:
98 """Checks that converting leaves tables that were converted alone.
100 Args:
101 data_folder: a data folder with the master files of 2004 and 2025.
103 """
104 paths = nibrs.convert(
105 2004, data_folder, tables = 'agencies', progress = False)
106 path = paths[2004]['agencies']
107 before = path.stat().st_mtime_ns
108 nibrs.convert(2004, data_folder, tables = 'agencies', progress = False)
109 assert path.stat().st_mtime_ns == before
110 assert not nibrs.store.table_path('victims', 2004, data_folder).exists()
113def test_load_tables(data_folder: pathlib.Path) -> None:
114 """Checks that several tables load as datasets.
116 Args:
117 data_folder: a data folder with the master files of 2004 and 2025.
119 """
120 datasets = nibrs.load_tables(
121 ['agencies', 'offenses'], 2004, folder = data_folder,
122 progress = False)
123 assert set(datasets) == {'agencies', 'offenses'}
124 assert all(isinstance(d, amos.Dataset) for d in datasets.values())
127def test_build_incidents(data_folder: pathlib.Path) -> None:
128 """Checks the built table of incidents.
130 Args:
131 data_folder: a data folder with the master files of 2004 and 2025.
133 """
134 built = tables.build(
135 'incidents', 2004, folder = data_folder, label = 'cleared',
136 groups = 'agency_ori')
137 assert isinstance(built, amos.Dataset)
138 assert (built.label, built.task, built.groups) == (
139 'cleared', 'classify', ['agency_ori'])
140 incidents = built.data
141 assert len(incidents) == 2
142 first = incidents.set_index('incident_number').loc['E0DROSC-2 N']
143 assert first['offense_count'] == 1
144 assert bool(first['against_property'])
145 assert bool(first['has_burglary_breaking_entering'])
146 assert not first['cleared']
147 assert first['agency_city'] == 'TEMPE'
148 assert first['agency_population'] == 158880
149 second = incidents.set_index('incident_number').loc['E0DROSCQ2 N']
150 assert bool(second['arrested'])
151 assert bool(second['cleared'])
152 assert second['days_to_arrest'] == 2
153 assert second['offender_count'] == 1
154 assert second['incident_weekday'] == 'Thursday'
157def test_build_victims_joins_offenders(data_folder: pathlib.Path) -> None:
158 """Checks that built victims have their offenders and agencies.
160 Args:
161 data_folder: a data folder with the master files of 2004 and 2025.
163 """
164 victims = tables.build('victims', 2004, folder = data_folder).data
165 person = victims.set_index('incident_number').loc['E0DROSCQ2 N']
166 assert person['relationship'] == 'Spouse'
167 assert person['offender_age'] == 31
168 assert person['offender_sex'] == 'Female'
169 assert person['agency_population'] == 158880
172def test_build_other_tables(data_folder: pathlib.Path) -> None:
173 """Checks the other built tables and unknown tables.
175 Args:
176 data_folder: a data folder with the master files of 2004 and 2025.
178 """
179 arrestees = tables.build('arrestees', 2004, folder = data_folder).data
180 assert arrestees.iloc[0]['days_to_arrest'] == 2
181 for name in ('offenses', 'offenders'):
182 built = tables.build(name, 2004, folder = data_folder)
183 assert 'agency_city' in built.data.columns
184 with pytest.raises(KeyError):
185 tables.build('suspects', 2004, folder = data_folder)
186 with pytest.raises(KeyError):
187 tables.incidents({})
190def test_decode_later(data_folder: pathlib.Path) -> None:
191 """Checks decoding a table that was loaded with codes.
193 Args:
194 data_folder: a data folder with the master files of 2004 and 2025.
196 """
197 offenses = nibrs.load(
198 'offenses', 2004, folder = data_folder, decode = False).data
199 assert set(offenses['offense_code']) == {'220', '13B'}
200 decoded = nibrs.decode(offenses.copy())
201 assert set(decoded['offense_code']) == {
202 'Burglary/Breaking & Entering', 'Simple Assault'}
203 custom = munge.decode(
204 offenses.copy(), ['location'], {'location': {'24': 'Shop'}})
205 assert 'Shop' in set(custom['location'])
206 assert set(custom['offense_code']) == {'220', '13B'}
209def test_unknown_codes_are_kept() -> None:
210 """Checks that codes that are not in a codebook are kept."""
211 data = pd.DataFrame({'location': pd.Categorical(['20', '99'])})
212 decoded = munge.decode(data)
213 assert list(decoded['location']) == ['Residence/Home', '99']
216def test_download_checks_years(
217 monkeypatch: pytest.MonkeyPatch,
218 tmp_path: pathlib.Path) -> None:
219 """Checks that downloading checks its years against the CDE's.
221 Args:
222 monkeypatch: replaces attributes for the test.
223 tmp_path: the test's temporary folder.
225 """
226 fetched: list[str] = []
227 monkeypatch.setattr(cde, '_year_range', lambda timeout: (1991, 2025))
228 monkeypatch.setattr(
229 cde, '_fetch',
230 lambda key, path, progress, timeout: fetched.append(key))
231 paths = cde.download('2024-2025', tmp_path, progress = False)
232 assert [p.name for p in paths] == ['nibrs-2024.zip', 'nibrs-2025.zip']
233 assert fetched == [
234 'master_files/nibrs/nibrs-2024.zip',
235 'master_files/nibrs/nibrs-2025.zip']
236 assert cde.available_years() == list(range(1991, 2026))
237 with pytest.raises(ValueError, match = '1990'):
238 cde.download(1990, tmp_path)
239 with pytest.raises(ValueError, match = 'no years'):
240 cde.download([], tmp_path)
243def test_download_saves_files(
244 monkeypatch: pytest.MonkeyPatch,
245 tmp_path: pathlib.Path) -> None:
246 """Checks that downloads are saved, without a connection.
248 Args:
249 monkeypatch: replaces attributes for the test.
250 tmp_path: the test's temporary folder.
252 """
254 class Response:
255 """A response of the CDE, without a connection.
257 Args:
258 payload: what `json` returns. Defaults to `None`.
260 """
262 """ Initialization Methods """
264 def __init__(self, payload: Any = None) -> None:
265 """Keeps the payload and says the file is 6 bytes long.
267 Args:
268 payload: what `json` returns. Defaults to `None`.
270 """
271 self.payload = payload
272 self.headers = {'Content-Length': '6'}
274 """ Public Methods """
276 def raise_for_status(self) -> None:
277 """Does nothing, since the request succeeded."""
278 return None
280 def json(self) -> Any:
281 """Returns the payload.
283 Returns:
284 The payload.
286 """
287 return self.payload
289 def iter_content(self, chunk_size: int) -> list[bytes]:
290 """Returns the blocks of the file.
292 Args:
293 chunk_size: the size of each block, which is not used.
295 Returns:
296 The file in two blocks.
298 """
299 return [b'abc', b'def']
301 """ Dunder Methods """
303 def __enter__(self) -> Response:
304 """Returns the response, for a `with` statement.
306 Returns:
307 The response.
309 """
310 return self
312 def __exit__(self, *args: Any) -> None:
313 """Does nothing at the end of a `with` statement.
315 Args:
316 *args: the exception, if any, which is not used.
318 """
319 return None
321 def get(url: str, **kwargs: Any) -> Response:
322 """Returns the CDE's response to a request.
324 Args:
325 url: the address of the request.
326 **kwargs: the request's parameters.
328 Returns:
329 A link for the signed URL service, the list of master files for
330 the CDE's list, and the file otherwise.
332 """
333 if url == cde._SIGNED_URL:
334 key = kwargs['params']['key']
335 return Response({key: 'https://example.test/file'})
336 if url == cde._MASTERS_URL:
337 return Response([{'id': 'nibrs', 'minYear': 1991, 'maxYear': 2025}])
338 return Response()
340 cde._year_range.cache_clear()
341 monkeypatch.setattr(cde.requests, 'get', get)
342 path = cde.download(2025, tmp_path, progress = True)[0]
343 assert path.read_bytes() == b'abcdef'
344 help_path = cde.download_documentation(tmp_path)
345 assert help_path.name == 'nibrs-help.zip'
346 cde._year_range.cache_clear()
349def test_download_resumes_with_new_links(
350 monkeypatch: pytest.MonkeyPatch,
351 tmp_path: pathlib.Path) -> None:
352 """Checks that a broken download resumes with a new link.
354 Args:
355 monkeypatch: replaces attributes for the test.
356 tmp_path: the test's temporary folder.
358 """
359 links: list[str] = []
360 ranges: list[str | None] = []
361 responses = iter([
362 _Broken([b'abc'], length = 9),
363 _Broken([b'def'], length = 6, status = 206),
364 _Broken([b'ghi'], length = 3, status = 206, breaks = False)])
366 def signed_url(key: str, timeout: float) -> str:
367 """Returns a new link each time, as the CDE does.
369 Args:
370 key: the key of the file, which is not used.
371 timeout: seconds to wait, which is not used.
373 Returns:
374 A link numbered by how many links came before it.
376 """
377 links.append(f'https://example.test/{len(links)}')
378 return links[-1]
380 def get(url: str, **kwargs: Any) -> _Broken:
381 """Returns the next response, noting the range that was asked for.
383 Args:
384 url: the address of the request, which is not used.
385 **kwargs: the request's parameters, with its headers.
387 Returns:
388 The next response.
390 """
391 ranges.append(kwargs['headers'].get('Range'))
392 return next(responses)
394 monkeypatch.setattr(cde, '_signed_url', signed_url)
395 monkeypatch.setattr(cde.requests, 'get', get)
396 path = tmp_path / 'nibrs-2025.zip'
397 cde._fetch('key', path, False, 1.0)
398 assert path.read_bytes() == b'abcdefghi'
399 assert len(links) == 3
400 assert ranges == [None, 'bytes=3-', 'bytes=6-']
401 assert not path.with_name('nibrs-2025.zip.part').exists()
404def test_download_starts_over_if_range_is_ignored(
405 monkeypatch: pytest.MonkeyPatch,
406 tmp_path: pathlib.Path) -> None:
407 """Checks that a download starts over if the server ignores its range.
409 Args:
410 monkeypatch: replaces attributes for the test.
411 tmp_path: the test's temporary folder.
413 """
414 responses = iter([
415 _Broken([b'abc'], length = 6),
416 _Broken([b'abcdef'], length = 6, breaks = False)])
417 monkeypatch.setattr(cde, '_signed_url', lambda key, timeout: 'link')
418 monkeypatch.setattr(
419 cde.requests, 'get', lambda url, **kwargs: next(responses))
420 path = tmp_path / 'nibrs-2025.zip'
421 cde._fetch('key', path, False, 1.0)
422 assert path.read_bytes() == b'abcdef'
425def test_download_gives_up_without_progress(
426 monkeypatch: pytest.MonkeyPatch,
427 tmp_path: pathlib.Path) -> None:
428 """Checks that a download gives up after tries that add nothing.
430 Args:
431 monkeypatch: replaces attributes for the test.
432 tmp_path: the test's temporary folder.
434 """
435 calls: list[str] = []
437 def get(url: str, **kwargs: Any) -> _Broken:
438 """Returns a response that downloads something only the first time.
440 Args:
441 url: the address of the request.
442 **kwargs: the request's parameters, which are not used.
444 Returns:
445 A response that breaks, with one block the first time and none
446 after.
448 """
449 calls.append(url)
450 return _Broken([b'abc'] if len(calls) == 1 else [], length = 9)
452 monkeypatch.setattr(cde, '_signed_url', lambda key, timeout: 'link')
453 monkeypatch.setattr(cde.requests, 'get', get)
454 path = tmp_path / 'nibrs-2025.zip'
455 with pytest.raises(cde.requests.exceptions.ChunkedEncodingError):
456 cde._fetch('key', path, False, 1.0)
457 # The first try downloaded something, so it does not count.
458 assert len(calls) == 1 + cde._ATTEMPTS
459 assert not path.exists()
462""" Private Classes """
465class _Broken:
466 """A response whose connection breaks after some blocks.
468 Args:
469 blocks: the blocks sent before the connection breaks.
470 length: the "Content-Length" of the response.
471 status: the status code. Defaults to 200.
472 breaks: whether the connection breaks after the blocks. Defaults to
473 `True`.
475 """
477 """ Initialization Methods """
479 def __init__(
480 self,
481 blocks: list[bytes],
482 length: int,
483 status: int = 200,
484 breaks: bool = True) -> None:
485 """Keeps the blocks, length, status, and whether the response breaks.
487 Args:
488 blocks: the blocks sent before the connection breaks.
489 length: the "Content-Length" of the response.
490 status: the status code. Defaults to 200.
491 breaks: whether the connection breaks after the blocks. Defaults
492 to `True`.
494 """
495 self.blocks = blocks
496 self.headers = {'Content-Length': str(length)}
497 self.status_code = status
498 self.breaks = breaks
500 """ Public Methods """
502 def raise_for_status(self) -> None:
503 """Does nothing, since the request succeeded."""
504 return None
506 def iter_content(self, chunk_size: int) -> Any:
507 """Yields the blocks, then breaks the connection if it breaks.
509 Args:
510 chunk_size: the size of each block, which is not used.
512 Yields:
513 The blocks.
515 Raises:
516 requests.exceptions.ChunkedEncodingError: after the blocks, if
517 the connection breaks.
519 """
520 yield from self.blocks
521 if self.breaks:
522 raise cde.requests.exceptions.ChunkedEncodingError('broken')
524 """ Dunder Methods """
526 def __enter__(self) -> _Broken:
527 """Returns the response, for a `with` statement.
529 Returns:
530 The response.
532 """
533 return self
535 def __exit__(self, *args: Any) -> None:
536 """Does nothing at the end of a `with` statement.
538 Args:
539 *args: the exception, if any, which is not used.
541 """
542 return None