diff options
Diffstat (limited to 'tests/incite/mergers/test_merge_collection.py')
| -rw-r--r-- | tests/incite/mergers/test_merge_collection.py | 73 |
1 files changed, 58 insertions, 15 deletions
diff --git a/tests/incite/mergers/test_merge_collection.py b/tests/incite/mergers/test_merge_collection.py index ec507bc..7ed3996 100644 --- a/tests/incite/mergers/test_merge_collection.py +++ b/tests/incite/mergers/test_merge_collection.py @@ -1,17 +1,22 @@ -from datetime import datetime, timezone, timedelta +from __future__ import annotations + +from datetime import UTC, datetime, timedelta from itertools import product +from typing import TYPE_CHECKING import pandas as pd import pytest from pandera.pandas import DataFrameSchema -from generalresearch.incite.mergers import ( +from generalresearch.incite.mergers.base import ( MergeCollection, MergeType, ) -from test_utils.incite.conftest import mnt_filepath -merge_types = list(e for e in MergeType if e != MergeType.TEST) +if TYPE_CHECKING: + from generalresearch.incite.base import GRLDatasets + +merge_types = [e for e in MergeType if e != MergeType.TEST] @pytest.mark.parametrize( @@ -21,17 +26,20 @@ merge_types = list(e for e in MergeType if e != MergeType.TEST) merge_types, ["5min", "6h", "14D"], [timedelta(days=30)], - [ - (datetime.now(tz=timezone.utc) - timedelta(days=35)).replace( - microsecond=0 - ) - ], + [(datetime.now(tz=UTC) - timedelta(days=35)).replace(microsecond=0)], ) ), ) class TestMergeCollection: - def test_init(self, mnt_filepath, merge_type, offset, duration, start): + def test_init( + self, + merge_type: MergeType, + offset: str, + duration: timedelta, + start: datetime, + mnt_filepath: GRLDatasets, + ): with pytest.raises(expected_exception=ValueError) as cm: MergeCollection(archive_path=mnt_filepath.data_src) assert "Must explicitly provide a merge_type" in str(cm.value) @@ -42,7 +50,14 @@ class TestMergeCollection: ) assert instance.merge_type == merge_type - def test_items(self, mnt_filepath, merge_type, offset, duration, start): + def test_items( + self, + mnt_filepath: GRLDatasets, + merge_type: MergeType, + offset: str, + duration: timedelta, + start: datetime, + ): instance = MergeCollection( merge_type=merge_type, offset=offset, @@ -53,7 +68,14 @@ class TestMergeCollection: assert len(instance.interval_range) == len(instance.items) - def test_progress(self, mnt_filepath, merge_type, offset, duration, start): + def test_progress( + self, + mnt_filepath: GRLDatasets, + merge_type: MergeType, + offset: str, + duration: timedelta, + start: datetime, + ): instance = MergeCollection( merge_type=merge_type, offset=offset, @@ -67,7 +89,14 @@ class TestMergeCollection: assert instance.progress.shape[1] == 7 assert instance.progress["group_by"].isnull().all() - def test_schema(self, mnt_filepath, merge_type, offset, duration, start): + def test_schema( + self, + offset: str, + duration: timedelta, + start: datetime, + mnt_filepath: GRLDatasets, + merge_type: MergeType, + ): instance = MergeCollection( merge_type=merge_type, archive_path=mnt_filepath.archive_path(enum_type=merge_type), @@ -75,7 +104,14 @@ class TestMergeCollection: assert isinstance(instance._schema, DataFrameSchema) - def test_load(self, mnt_filepath, merge_type, offset, duration, start): + def test_load( + self, + mnt_filepath: GRLDatasets, + merge_type: MergeType, + offset: str, + duration: timedelta, + start: datetime, + ): instance = MergeCollection( merge_type=merge_type, start=start, @@ -87,7 +123,14 @@ class TestMergeCollection: # Confirm that there are no archives available yet assert instance.progress.has_archive.eq(False).all() - def test_get_items(self, mnt_filepath, merge_type, offset, duration, start): + def test_get_items( + self, + mnt_filepath: GRLDatasets, + merge_type: MergeType, + offset: str, + duration: timedelta, + start: datetime, + ): instance = MergeCollection( start=start, finished=start + duration, |