Source OneSignal: correct get_update_state and unit test

airbyte-integrations/connectors/source-onesignal/source_onesignal/streams.py

@@ -4,10 +4,9 @@
 
 import time
 from abc import ABC
-from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, TypeVar
+from typing import Any, Iterable, List, Mapping, MutableMapping, Optional
 
 import pendulum
-import pydantic
 import requests
 from airbyte_cdk.models import SyncMode
 from airbyte_cdk.sources.streams.http import HttpStream, HttpSubStream
@@ -76,28 +75,6 @@ def parse_response(self, response: requests.Response, **kwargs) -> Iterable[Mapp
         yield from data
 
 
-T = TypeVar("T")
-
-
-class StateValueWrapper(pydantic.BaseModel):
-    stream: T
-    state_value: int
-    max_cursor_time = 0
-
-    def __repr__(self):
-        """Overrides print view"""
-        return str(self.value)
-
-    @property
-    def value(self):
-        """Return max cursor time after stream sync is finished."""
-        return self.max_cursor_time if self.stream.is_finished else self.state_value
-
-    def dict(self, **kwargs):
-        """Overrides default logic to return current value only."""
-        return {pydantic.utils.ROOT_KEY: self.value}
-
-
 class ChildStreamMixin(HttpSubStream):
 
     is_finished = False
@@ -176,14 +153,11 @@ def get_updated_state(
         current_stream_state: MutableMapping[str, Any],
         latest_record: Mapping[str, Any],
     ) -> Mapping[str, Any]:
-        state_value = (current_stream_state or {}).get(self.cursor_field, 0)
-        if not isinstance(state_value, StateValueWrapper):
-            state_value = StateValueWrapper(stream=self, state_value=state_value)
-
-        record_time = latest_record.get(self.cursor_field, self.start_date)
-        state_value.max_cursor_time = max(state_value.max_cursor_time, record_time)
+        current_stream_state = current_stream_state or {}
+        current_stream_state_date = current_stream_state.get(self.cursor_field, self.start_date)
+        latest_record_date = latest_record.get(self.cursor_field, self.start_date)
 
-        return {self.cursor_field: state_value}
+        return {self.cursor_field: max(current_stream_state_date, latest_record_date)}
 
 
 class Apps(OnesignalStream):

airbyte-integrations/connectors/source-onesignal/unit_tests/test_incremental_streams.py

@@ -41,21 +41,21 @@ def test_cursor_field(stream):
 
 def test_get_updated_state(stream):
     inputs = {"current_stream_state": {"updated_at": 42}, "latest_record": {"updated_at": 90}}
-    expected_state = 42
+    expected_state = 90
     state = stream.get_updated_state(**inputs)
-    assert state["updated_at"].value == expected_state
+    assert state["updated_at"] == expected_state
 
     inputs = {"current_stream_state": state, "latest_record": {"updated_at": 100}}
-    expected_state = 42
+    expected_state = 100
     state = stream.get_updated_state(**inputs)
-    assert state["updated_at"].value == expected_state
+    assert state["updated_at"] == expected_state
 
     # after stream sync is finished, state should output the max cursor time
     stream.is_finished = True
     inputs = {"current_stream_state": state, "latest_record": {"updated_at": 80}}
     expected_state = 100
     state = stream.get_updated_state(**inputs)
-    assert state["updated_at"].value == expected_state
+    assert state["updated_at"] == expected_state
 
 
 def test_stream_slices(stream, requests_mock):
@@ -94,9 +94,9 @@ def test_end_of_stream_state(parent, args, requests_mock):
         for record in stream.read_records(sync_mode, stream_slice=app_slice):
             state = stream.get_updated_state(state, record)
             if idx == 2:  # the last slice
-                assert state["queued_at"].value == 90
+                assert state["queued_at"] == 90
             else:
-                assert state["queued_at"].value == 50
+                assert state["queued_at"] == 90
 
 
 def test_supports_incremental(patch_incremental_base_class, mocker, parent, args):