fix tests

airbytehq · Jun 15, 2024 · 207dac1 · 207dac1
1 parent c2a211e
commit 207dac1
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 11 deletions.
diff --git a/airbyte-integrations/connectors/source-github/source_github/streams.py b/airbyte-integrations/connectors/source-github/source_github/streams.py
@@ -48,6 +48,7 @@ def __init__(self, api_url: str = "https://api.github.com", access_token_type: s
 
         self.access_token_type = access_token_type
         self.api_url = api_url
+        self.state = {}
 
     @property
     def url_base(self) -> str:
@@ -1340,6 +1341,7 @@ def read_records(
             sync_mode=sync_mode, cursor_field=cursor_field, stream_slice=stream_slice, stream_state=stream_state
         ):
             if not starting_point or record[self.cursor_field] > starting_point:
+                self.state = self._get_updated_state(self.state, record)
                 yield record
 
     def get_starting_point(self, stream_state: Mapping[str, Any], stream_slice: Mapping[str, Any]) -> str:
@@ -1353,7 +1355,7 @@ def get_starting_point(self, stream_state: Mapping[str, Any], stream_slice: Mapp
                 return stream_state_value
         return self._start_date
 
-    def get_updated_state(self, current_stream_state: MutableMapping[str, Any], latest_record: Mapping[str, Any]):
+    def _get_updated_state(self, current_stream_state: MutableMapping[str, Any], latest_record: Mapping[str, Any]):
         repository = latest_record["repository"]
         project_id = str(latest_record["project_id"])
         updated_state = latest_record[self.cursor_field]
@@ -1410,6 +1412,7 @@ def read_records(
             sync_mode=sync_mode, cursor_field=cursor_field, stream_slice=stream_slice, stream_state=stream_state
         ):
             if not starting_point or record[self.cursor_field] > starting_point:
+                self.state = self._get_updated_state(self.state, record)
                 yield record
 
     def get_starting_point(self, stream_state: Mapping[str, Any], stream_slice: Mapping[str, Any]) -> str:
@@ -1424,7 +1427,7 @@ def get_starting_point(self, stream_state: Mapping[str, Any], stream_slice: Mapp
                 return stream_state_value
         return self._start_date
 
-    def get_updated_state(self, current_stream_state: MutableMapping[str, Any], latest_record: Mapping[str, Any]):
+    def _get_updated_state(self, current_stream_state: MutableMapping[str, Any], latest_record: Mapping[str, Any]):
         repository = latest_record["repository"]
         project_id = str(latest_record["project_id"])
         column_id = str(latest_record["column_id"])

diff --git a/airbyte-integrations/connectors/source-github/unit_tests/test_migrations/test_config.json b/airbyte-integrations/connectors/source-github/unit_tests/test_migrations/test_config.json
@@ -1,8 +1 @@
-{
-  "credentials": {
-    "personal_access_token": "personal_access_token"
-  },
-  "repository": "airbytehq/airbyte airbytehq/airbyte-platform",
-  "start_date": "2000-01-01T00:00:00Z",
-  "branch": "airbytehq/airbyte/master airbytehq/airbyte-platform/main"
-}
+{"credentials": {"personal_access_token": "personal_access_token"}, "repository": "airbytehq/airbyte airbytehq/airbyte-platform", "start_date": "2000-01-01T00:00:00Z", "branch": "airbytehq/airbyte/master airbytehq/airbyte-platform/main"}
diff --git a/airbyte-integrations/connectors/source-github/unit_tests/utils.py b/airbyte-integrations/connectors/source-github/unit_tests/utils.py
@@ -19,7 +19,7 @@ def read_incremental(stream_instance: Stream, stream_state: MutableMapping[str,
     for slice in slices:
         records = stream_instance.read_records(sync_mode=SyncMode.incremental, stream_slice=slice, stream_state=stream_state)
         for record in records:
-            stream_state = stream_instance.get_updated_state(stream_state, record)
+            stream_state = stream_instance._get_updated_state(stream_state, record)
             res.append(record)
     return res