Skip to content
Snippets Groups Projects
conftest.py 2.13 KiB
Newer Older
  • Learn to ignore specific revisions
  • import numpy
    from obspy import Stream, UTCDateTime
    from obspy.clients.neic.client import Client
    import pytest
    
    from geomagio import TimeseriesUtility
    
    
    @pytest.fixture(scope="class")
    def MockMiniSeedClient() -> Client:
        """replaces default obspy miniseed client's get_waveforms method to return trace of ones"""
    
        class MockMiniSeedClient(Client):
            def get_waveforms(
                self,
                network: str,
                station: str,
                location: str,
                channel: str,
                starttime: UTCDateTime,
                endtime: UTCDateTime,
            ):
                trace = TimeseriesUtility.create_empty_trace(
                    starttime=starttime,
                    endtime=endtime,
                    observatory=station,
                    channel=channel,
                    type=self._get_data_type(location=location),
                    interval=self._get_interval(channel=channel),
                    network=network,
                    station=station,
                    location=location,
                )
                trace.data = numpy.ones(trace.stats.npts)
                return Stream([trace])
    
            def _get_interval(self, channel: str) -> str:
                channel_start = channel[0]
                if channel_start == "B":
                    return "tenhertz"
                elif channel_start == "L":
                    return "second"
                elif channel_start == "U":
                    return "minute"
                elif channel_start == "R":
                    return "hour"
                elif channel_start == "P":
                    return "day"
                else:
                    raise ValueError(f"Unexpected channel start: {channel_start}")
    
            def _get_data_type(self, location: str) -> str:
                location_start = location[0]
                if location_start == "R":
                    return "variation"
                elif location_start == "A":
                    return "adjusted"
                elif location_start == "Q":
                    return "quasi-definitive"
                elif location_start == "D":
                    return "definitive"
                else:
                    raise ValueError(f"Unexpected location start: {location_start}")
    
        yield MockMiniSeedClient