fix: pr review

apify · Feb 7, 2023 · 6d9862f · 6d9862f
1 parent 554ea28
commit 6d9862f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/tests/integration/test_actor_events.py b/tests/integration/test_actor_events.py
@@ -2,14 +2,13 @@
 
 from apify import Actor
 from apify.consts import ActorEventType
-from apify_client import ApifyClientAsync
 
 from .conftest import ActorFactory
 
 
 class TestActorEvents:
 
-    async def test_interval_events(self, make_actor: ActorFactory, apify_client_async: ApifyClientAsync) -> None:
+    async def test_interval_events(self, make_actor: ActorFactory) -> None:
         async def main() -> None:
             import os
 
@@ -33,7 +32,7 @@ async def log_event(data):  # type: ignore
 
         assert run_result is not None
         assert run_result['status'] == 'SUCCEEDED'
-        dataset_items_page = await apify_client_async.dataset(run_result['defaultDatasetId']).list_items()
+        dataset_items_page = await actor.last_run().dataset().list_items()
         persist_state_events = [item for item in dataset_items_page.items if item['event_type'] == ActorEventType.PERSIST_STATE]
         system_info_events = [item for item in dataset_items_page.items if item['event_type'] == ActorEventType.SYSTEM_INFO]
         assert len(persist_state_events) > 2

diff --git a/tests/integration/test_actor_lifecycle.py b/tests/integration/test_actor_lifecycle.py
@@ -1,5 +1,4 @@
 from apify import Actor
-from apify_client import ApifyClientAsync
 
 from .conftest import ActorFactory
 
@@ -103,7 +102,7 @@ async def main() -> None:
 
 class TestActorMain:
 
-    async def test_actor_main(self, make_actor: ActorFactory, apify_client_async: ApifyClientAsync) -> None:
+    async def test_actor_main(self, make_actor: ActorFactory) -> None:
         async def main() -> None:
             async def actor_function() -> None:
                 input = await Actor.get_input()
@@ -141,6 +140,6 @@ async def actor_function() -> None:
         run_with_output = await actor.call(run_input={'set_output': test_output})
         assert run_with_output is not None
         assert run_with_output['status'] == 'SUCCEEDED'
-        output = await apify_client_async.key_value_store(run_with_output['defaultKeyValueStoreId']).get_record('OUTPUT')
+        output = await actor.last_run().key_value_store().get_record('OUTPUT')
         assert output is not None
         assert output['value'] == test_output