Improved tests (added cleanup)

instagrapi/extractors.py

@@ -45,7 +45,7 @@ def extract_media_v1(data):
         "product_type": product_type,
         "code": data["code"],
         "thumbnail_url": thumbnail_url,
-        "location": location,
+        "location": location or {},
         "user": extract_user_short(user),
         "comment_count": int(data.get("comment_count") or 0),
         # the media just published has no like_count
@@ -113,7 +113,7 @@ def extract_media_gql(data):
             data.get("display_resources", data.get('thumbnail_resources')),
             key=lambda o: o["config_width"] * o["config_height"],
         ).pop()["src"],
-        "location": location,
+        "location": location or {},
         "user": user,
         "comment_count": json_value(data, "edge_media_to_comment", "count"),
         "like_count": json_value(data, "edge_media_preview_like", "count"),

tests.py

@@ -12,6 +12,15 @@
 ACCOUNT_PASSWORD = os.environ.get("IG_PASSWORD", "yoa5af6deeRujeec")
 
 
+def cleanup(*paths):
+    for path in paths:
+        try:
+            os.remove(path)
+            os.remove(f'{path}.jpg')
+        except FileNotFoundError:
+            continue
+
+
 class FakeClientTestCase(unittest.TestCase):
     api = None
 
@@ -173,7 +182,7 @@ def test_media_edit(self):
             self.assertEqual(media["caption_text"], msg)
             self.assertTrue(self.api.media_delete(media_pk))
         finally:
-            os.remove(path)
+            cleanup(path)
 
     def test_media_edit_igtv(self):
         media_pk = self.api.media_pk_from_url(
@@ -205,7 +214,7 @@ def test_media_edit_igtv(self):
             self.assertEqual(media["caption_text"], msg)
             self.assertTrue(self.api.media_delete(media["id"]))
         finally:
-            os.remove(path)
+            cleanup(path)
 
     def test_media_user(self):
         user = self.api.media_user(2154602296692269830)
@@ -305,6 +314,9 @@ def test_two_extract_media_photo(self):
         self.assertTrue(media_v1.pop("thumbnail_url").startswith("https://"))
         self.assertTrue(media_gql.pop("thumbnail_url").startswith("https://"))
         self.assertTrue(media_v1.pop("comment_count") <= media_gql.pop("comment_count"))
+        location_v1, location_gql = media_v1.pop('location'), media_gql.pop('location')
+        for key in ('name', 'pk'):
+            self.assertEqual(location_v1.get(key), location_gql.get(key))
         self.assertDictEqual(media_v1, media_gql)
 
     def test_two_extract_media_video(self):
@@ -316,6 +328,9 @@ def test_two_extract_media_video(self):
         self.assertTrue(media_v1.pop("video_url").startswith("https://"))
         self.assertTrue(media_gql.pop("thumbnail_url").startswith("https://"))
         self.assertTrue(media_gql.pop("video_url").startswith("https://"))
+        location_v1, location_gql = media_v1.pop('location'), media_gql.pop('location')
+        for key in ('name', 'pk'):
+            self.assertEqual(location_v1.get(key), location_gql.get(key))
         self.assertDictEqual(media_v1, media_gql)
 
     def test_two_extract_media_album(self):
@@ -331,6 +346,9 @@ def test_two_extract_media_album(self):
             self.assertTrue(res.pop("thumbnail_url").startswith("https://"))
             if res['media_type'] == 2:
                 self.assertTrue(res.pop("video_url").startswith("https://"))
+        location_v1, location_gql = media_v1.pop('location'), media_gql.pop('location')
+        for key in ('name', 'pk'):
+            self.assertEqual(location_v1.get(key), location_gql.get(key))
         self.assertDictEqual(media_v1, media_gql)
 
     def test_two_extract_media_igtv(self):
@@ -342,6 +360,9 @@ def test_two_extract_media_igtv(self):
         self.assertTrue(media_v1.pop("video_url").startswith("https://"))
         self.assertTrue(media_gql.pop("thumbnail_url").startswith("https://"))
         self.assertTrue(media_gql.pop("video_url").startswith("https://"))
+        location_v1, location_gql = media_v1.pop('location'), media_gql.pop('location')
+        for key in ('name', 'pk'):
+            self.assertEqual(location_v1.get(key), location_gql.get(key))
         self.assertDictEqual(media_v1, media_gql)
 
 
@@ -474,7 +495,7 @@ def test_photo_upload_without_location(self):
             self.assertEqual(media["caption_text"], "Test caption for photo")
             self.assertFalse(media["location"])
         finally:
-            os.remove(path)
+            cleanup(path)
             self.assertTrue(self.api.media_delete(media["id"]))
 
     def test_photo_upload(self):
@@ -492,7 +513,7 @@ def test_photo_upload(self):
             for key, val in {'pk': 213597007, 'name': 'Palace Square', 'lat': 59.939166666667, 'lng': 30.315833333333}.items():
                 self.assertEqual(media["location"][key], val)
         finally:
-            os.remove(path)
+            cleanup(path)
             self.assertTrue(self.api.media_delete(media["id"]))
 
     def test_video_upload(self):
@@ -510,7 +531,7 @@ def test_video_upload(self):
             for key, val in {'pk': 213597007, 'name': 'Palace Square', 'lat': 59.939166666667, 'lng': 30.315833333333}.items():
                 self.assertEqual(media["location"][key], val)
         finally:
-            os.remove(path)
+            cleanup(path)
             self.assertTrue(self.api.media_delete(media["id"]))
 
     def test_album_upload(self):
@@ -527,8 +548,7 @@ def test_album_upload(self):
             for key, val in {'pk': 213597007, 'name': 'Palace Square', 'lat': 59.939166666667, 'lng': 30.315833333333}.items():
                 self.assertEqual(media["location"][key], val)
         finally:
-            for path in paths:
-                os.remove(path)
+            cleanup(*paths)
             self.assertTrue(self.api.media_delete(media["id"]))
 
     def test_igtv_upload(self):
@@ -548,7 +568,7 @@ def test_igtv_upload(self):
             for key, val in {'pk': 213597007, 'name': 'Palace Square', 'lat': 59.939166666667, 'lng': 30.315833333333}.items():
                 self.assertEqual(media["location"][key], val)
         finally:
-            os.remove(path)
+            cleanup(path)
             self.assertTrue(self.api.media_delete(media["id"]))