Repository: PrefectHQ/prefect-twitter (language: python)

Issue:
Create get_media_upload_status task
https://docs.tweepy.org/en/stable/api.html#tweepy.API.get_media_upload_status get Media object with the Media ID
Update .pre-commit-config.yaml



Discussion / hints:

Fix patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index e09fb4c..47b4b4b 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -5,7 +5,7 @@ repos:
       - id: isort
         language_version: python3
   - repo: https://github.com/psf/black
-    rev: 22.1.0
+    rev: 22.3.0
     hooks:
       - id: black
         language_version: python3
@@ -18,4 +18,4 @@ repos:
     hooks:
       - id: interrogate
         args: [-vv]
-        pass_filenames: false
\ No newline at end of file
+        pass_filenames: false
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 33d8c73..c0448f9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,4 +25,6 @@ Released on ????? ?th, 20??.
 
 ### Added
 
-- `task_name` task - [#1](https://github.com/PrefectHQ/prefect-twitter/pull/1)
+- `update_status` and `media_upload` tasks - [#1](https://github.com/PrefectHQ/prefect-twitter/pull/1)
+- `get_status` and `get_media_upload_status` tasks - [#2](https://github.com/PrefectHQ/prefect-twitter/pull/2)
+
diff --git a/prefect_twitter/media.py b/prefect_twitter/media.py
index 1824491..5d4b9db 100644
--- a/prefect_twitter/media.py
+++ b/prefect_twitter/media.py
@@ -10,6 +10,8 @@
 if TYPE_CHECKING:
     from io import IOBase
 
+    from tweepy import Media
+
     from prefect_twitter import TwitterCredentials
 
 
@@ -28,7 +30,7 @@ async def media_upload(
     Args:
         filename: The filename of the image to upload.
             This field is used for MIME type detection.
-        twitter_credentials: Credentials to use for authentication with Tweepy.
+        twitter_credentials: Credentials to use for authentication with Twitter.
         file: A file object to upload. If not specified, this task will attempt to
             locate and upload a file with the name specified in filename.
         chunked: Whether or not to use chunked media upload.
@@ -69,3 +71,46 @@ def example_media_upload_flow():
     media = await to_thread.run_sync(partial_media)
     media_id = media.media_id
     return media_id
+
+
+@task
+async def get_media_upload_status(
+    media_id: int, twitter_credentials: "TwitterCredentials"
+) -> "Media":
+    """
+    Check on the progress of a chunked media upload. If the upload has succeeded,
+    it's safe to create a Tweet with this media_id.
+
+    Args:
+        media_id: The ID of the media to check.
+        twitter_credentials: Credentials to use for authentication with Twitter.
+
+    Returns:
+        The Media object.
+
+    Example:
+        Tweets an update with just text.
+        ```python
+        from prefect import flow
+        from prefect_twitter import TwitterCredentials
+        from prefect_twitter.media import get_media_upload_status
+
+        @flow
+        def example_get_media_upload_status_flow():
+            twitter_credentials = TwitterCredentials(
+                consumer_key=consumer_key,
+                consumer_secret=consumer_secret,
+                access_token=access_token,
+                access_token_secret=access_token_secret
+            )
+            media_id = 1443668738906234883
+            media = get_media_upload_status(media_id, twitter_credentials)
+            return media
+
+        example_get_media_upload_status_flow()
+        ```
+    """
+    api = twitter_credentials.get_api()
+    partial_get = partial(api.get_media_upload_status, media_id)
+    media = await to_thread.run_sync(partial_get)
+    return media
diff --git a/prefect_twitter/tweets.py b/prefect_twitter/tweets.py
index cae54f7..42d7df4 100644
--- a/prefect_twitter/tweets.py
+++ b/prefect_twitter/tweets.py
@@ -7,6 +7,8 @@
 from prefect import get_run_logger, task
 
 if TYPE_CHECKING:
+    from tweepy import Status
+
     from prefect_twitter import TwitterCredentials
 
 
@@ -21,7 +23,7 @@ async def update_status(
     Updates the authenticating user's current status, also known as Tweeting.
 
     Args:
-        twitter_credentials: Credentials to use for authentication with Tweepy.
+        twitter_credentials: Credentials to use for authentication with Twitter.
         status: Text of the Tweet being created. This field is required
             if media_ids is not present.
         media_ids: A list of Media IDs being attached to the Tweet.
@@ -90,3 +92,46 @@ def example_update_status_flow():
     )
     status = await to_thread.run_sync(partial_update)
     return status.id
+
+
+@task
+async def get_status(
+    status_id: int, twitter_credentials: "TwitterCredentials", **kwargs: dict
+) -> "Status":
+    """
+    Returns a single status specified by the ID parameter.
+
+    Args:
+        status_id: The ID of the status.
+        twitter_credentials: Credentials to use for authentication with Twitter.
+        kwargs: Additional keyword arguments to pass to
+            [get_status](https://docs.tweepy.org/en/stable/api.html#tweepy.API.get_status).
+    Returns:
+        The Status object.
+
+    Example:
+        Tweets an update with just text.
+        ```python
+        from prefect import flow
+        from prefect_twitter import TwitterCredentials
+        from prefect_twitter.tweets import get_status
+
+        @flow
+        def example_get_status_flow():
+            twitter_credentials = TwitterCredentials(
+                consumer_key=consumer_key,
+                consumer_secret=consumer_secret,
+                access_token=access_token,
+                access_token_secret=access_token_secret
+            )
+            status_id = 1504591031626571777
+            status = get_status(status_id, twitter_credentials)
+            return status
+
+        example_get_status_flow()
+        ```
+    """
+    api = twitter_credentials.get_api()
+    partial_get = partial(api.get_status, status_id, **kwargs)
+    status = await to_thread.run_sync(partial_get)
+    return status


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 8a5ebb8..afa5c48 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -13,6 +13,12 @@ def update_status(self, status=None, media_ids=None, **kwargs):
         else:
             return MagicMock(id=status)
 
+    def get_status(self, status_id=None, **kwargs):
+        return status_id
+
+    def get_media_upload_status(self, media_id=None, **kwargs):
+        return media_id
+
 
 @pytest.fixture
 def twitter_credentials(monkeypatch):
diff --git a/tests/test_media.py b/tests/test_media.py
index 7312e54..04d3535 100644
--- a/tests/test_media.py
+++ b/tests/test_media.py
@@ -1,6 +1,6 @@
 from prefect import flow
 
-from prefect_twitter.media import media_upload
+from prefect_twitter.media import get_media_upload_status, media_upload
 
 
 def test_media_upload(twitter_credentials):
@@ -14,3 +14,14 @@ def test_flow():
     result = test_flow().result().result()
     assert result["media_id"] == path
     assert result["chunked"] is False
+
+
+def test_get_media_upload_status(twitter_credentials):
+    media_id = 42
+
+    @flow
+    def test_flow():
+        media = get_media_upload_status(media_id, twitter_credentials)
+        return media
+
+    assert test_flow().result().result() == media_id
diff --git a/tests/test_tweets.py b/tests/test_tweets.py
index 0f7fdd0..da459c4 100644
--- a/tests/test_tweets.py
+++ b/tests/test_tweets.py
@@ -1,7 +1,7 @@
 import pytest
 from prefect import flow
 
-from prefect_twitter.tweets import update_status
+from prefect_twitter.tweets import get_status, update_status
 
 
 def test_update_status(twitter_credentials):
@@ -9,8 +9,8 @@ def test_update_status(twitter_credentials):
 
     @flow
     def test_flow():
-        media = update_status(twitter_credentials, status=status)
-        return media
+        status_id = update_status(twitter_credentials, status=status)
+        return status_id
 
     assert test_flow().result().result() == status
 
@@ -18,8 +18,8 @@ def test_flow():
 def test_update_status_missing(twitter_credentials):
     @flow
     def test_flow():
-        media = update_status(twitter_credentials)
-        return media
+        status_id = update_status(twitter_credentials)
+        return status_id
 
     with pytest.raises(ValueError, match="One of text or media_ids"):
         test_flow().result().result()
@@ -30,7 +30,18 @@ def test_update_status_media(twitter_credentials):
 
     @flow
     def test_flow():
-        media = update_status(twitter_credentials, media_ids=media_ids)
-        return media
+        status_id = update_status(twitter_credentials, media_ids=media_ids)
+        return status_id
 
     assert test_flow().result().result() == media_ids[0]
+
+
+def test_get_status(twitter_credentials):
+    status_id = 42
+
+    @flow
+    def test_flow():
+        status = get_status(status_id, twitter_credentials)
+        return status
+
+    assert test_flow().result().result() == status_id