47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
#
|
|
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
|
|
#
|
|
|
|
import datetime
|
|
|
|
import pytest
|
|
from airbyte_cdk.models import (
|
|
AirbyteMessage,
|
|
AirbyteStreamStatus,
|
|
AirbyteStreamStatusTraceMessage,
|
|
AirbyteTraceMessage,
|
|
StreamDescriptor,
|
|
TraceType,
|
|
)
|
|
from airbyte_cdk.models import Type as MessageType
|
|
|
|
|
|
@pytest.fixture
|
|
def remove_stack_trace():
|
|
def _remove_stack_trace(message: AirbyteMessage) -> AirbyteMessage:
|
|
"""
|
|
Helper method that removes the stack trace from Airbyte trace messages to make asserting against expected records easier
|
|
"""
|
|
if message.trace and message.trace.error and message.trace.error.stack_trace:
|
|
message.trace.error.stack_trace = None
|
|
return message
|
|
|
|
return _remove_stack_trace
|
|
|
|
|
|
@pytest.fixture
|
|
def as_stream_status():
|
|
def _as_stream_status(stream: str, status: AirbyteStreamStatus) -> AirbyteMessage:
|
|
trace_message = AirbyteTraceMessage(
|
|
emitted_at=datetime.datetime.now().timestamp() * 1000.0,
|
|
type=TraceType.STREAM_STATUS,
|
|
stream_status=AirbyteStreamStatusTraceMessage(
|
|
stream_descriptor=StreamDescriptor(name=stream),
|
|
status=status,
|
|
),
|
|
)
|
|
|
|
return AirbyteMessage(type=MessageType.TRACE, trace=trace_message)
|
|
|
|
return _as_stream_status
|