Refactor updater

cmd/updater/dynamo.py

@@ -0,0 +1,110 @@
+import boto3
+import os
+import time
+from updater import DEFAULT_PAGE_SIZE, _get_updates, _get_format
+from datetime import datetime
+
+dynamodb = boto3.resource('dynamodb')
+
+feeds_table = dynamodb.Table(os.getenv('UPDATER_DYNAMO_FEEDS_TABLE', 'Feeds'))
+
+
+def _update_feed(hash_id):
+    print('Updating feed {}'.format(hash_id))
+    feed = _query_feed(hash_id)
+
+    page_size = int(feed.get('PageSize', DEFAULT_PAGE_SIZE))
+    last_id = feed.get('LastID', None)
+    episodes = feed.get('Episodes', [])
+    item_url = feed['ItemURL']
+
+    # Format parameters
+    fmt = feed.get('Format', 'video')
+    quality = feed.get('Quality', 'high')
+
+    # Rebuild episode list from scratch
+    if not last_id:
+        episodes = []
+
+    start = time.time()
+    _, items, new_last_id = _get_updates(1, page_size, item_url, _get_format(fmt, quality), last_id)
+    end = time.time()
+
+    print('Got feed update: new {}, current {}. Update took: {}'.format(len(items), len(episodes), end-start))
+
+    # Update feed and submit back to Dynamo
+
+    unix_time = int(datetime.utcnow().timestamp())
+    feed['UpdatedAt'] = unix_time
+
+    if len(items) > 0:
+        episodes = items + episodes  # Prepand new episodes
+        del episodes[page_size:]  # Truncate list
+        feed['Episodes'] = episodes
+
+        # Update last seen video ID
+        feed['LastID'] = new_last_id
+
+        _update_feed_episodes(hash_id, feed)
+    else:
+        # Update last access field only
+        _update_feed_updated_at(hash_id, unix_time)
+
+
+def _query_feed(hash_id):
+    response = feeds_table.get_item(
+        Key={'HashID': hash_id},
+        ProjectionExpression='#prov,#type,#size,#fmt,#quality,#level,#id,#last_id,#episodes,#updated_at,#item_url',
+        ExpressionAttributeNames={
+            '#prov': 'Provider',
+            '#type': 'LinkType',
+            '#size': 'PageSize',
+            '#fmt': 'Format',
+            '#quality': 'Quality',
+            '#level': 'FeatureLevel',
+            '#id': 'ItemID',
+            '#last_id': 'LastID',
+            '#episodes': 'Episodes',
+            '#updated_at': 'UpdatedAt',
+            '#item_url': 'ItemURL',
+        },
+    )
+
+    item = response['Item']
+    return item
+
+
+def _update_feed_episodes(hash_id, feed):
+    feeds_table.update_item(
+        Key={
+            'HashID': hash_id,
+        },
+        UpdateExpression='SET #updated_at = :updated_at, #episodes = :episodes, #last_id = :last_id',
+        ExpressionAttributeNames={
+            '#updated_at': 'UpdatedAt',
+            '#episodes': 'Episodes',
+            '#last_id': 'LastID',
+        },
+        ExpressionAttributeValues={
+            ':updated_at': feed['UpdatedAt'],
+            ':episodes': feed['Episodes'],
+            ':last_id': feed['LastID'],
+        },
+        ReturnValues='NONE',
+    )
+
+
+def _update_feed_updated_at(hash_id, updated_at):
+    feeds_table.update_item(
+        Key={
+            'HashID': hash_id,
+        },
+        UpdateExpression='SET #updated_at = :updated_at',
+        ExpressionAttributeNames={
+            '#updated_at': 'UpdatedAt',
+        },
+        ExpressionAttributeValues={
+            ':updated_at': updated_at,
+        },
+        ReturnValues='NONE',
+    )

cmd/updater/dynamo_test.py

@@ -0,0 +1,10 @@
+import dynamo
+import unittest
+
+TEST_FEED = '86qZ'
+
+
+class TestUpdater(unittest.TestCase):
+    @unittest.skip
+    def test_update_feed(self):
+        dynamo._update_feed(TEST_FEED)

cmd/updater/function.py

@@ -0,0 +1,24 @@
+from updater import DEFAULT_PAGE_SIZE, _get_updates, _get_format
+
+
+# AWS Lambda entry point
+def handler(event, context):
+    url = event.get('url', None)
+    start = event.get('start', 1)
+    count = event.get('count', DEFAULT_PAGE_SIZE)
+
+    # Last seen video ID
+    last_id = event.get('last_id', None)
+
+    # Detect item format
+    fmt = event.get('format', 'video')
+    quality = event.get('quality', 'high')
+    ytdl_fmt = _get_format(fmt, quality)
+
+    print('Getting updates for %s (start=%d, count=%d, fmt: %s, last id: %s)' % (url, start, count, ytdl_fmt, last_id))
+    _, episodes, new_last_id = _get_updates(start, count, url, ytdl_fmt, last_id)
+
+    return {
+        'LastID': new_last_id,
+        'Episodes': episodes,
+    }

cmd/updater/function_test.py

@@ -0,0 +1,18 @@
+import function
+import unittest
+
+TEST_URL = 'https://www.youtube.com/user/CNN/videos'
+
+
+class TestUpdater(unittest.TestCase):
+    #@unittest.skip('heavy test, run manually')
+    def test_get_50(self):
+        resp = function.handler({
+            'url': 'https://www.youtube.com/channel/UCd6MoB9NC6uYN2grvUNT-Zg',
+            'start': 1,
+            'count': 50,
+            'format': 'audio',
+            'quality': 'low',
+        }, None)
+        self.assertEqual(len(resp['Episodes']), 50)
+        self.assertEqual(resp['Episodes'][0]['ID'], resp['LastID'])

cmd/updater/updater.py

@@ -1,133 +1,11 @@
 import youtube_dl
-import boto3
-import os
-import time
 from datetime import datetime
 
 BEST_FORMAT = "bestvideo+bestaudio/best"
 DEFAULT_PAGE_SIZE = 50
 
-dynamodb = boto3.resource('dynamodb')
-
-feeds_table = dynamodb.Table(os.getenv('UPDATER_DYNAMO_FEEDS_TABLE', 'Feeds'))
-
-
-def handler(event, context):
-    url = event.get('url', None)
-    if not url:
-        raise ValueError('Invalid resource URL %s' % url)
-
-    start = event.get('start', 1)
-    count = event.get('count', DEFAULT_PAGE_SIZE)
-
-    kind = event.get('kind', 'video_high')
-    last_id = event.get('last_id', None)
-
-    print('Getting updated for %s (start=%d, count=%d, kind: %s, last id: %s)' % (url, start, count, kind, last_id))
-    return _get_updates(start, count, url, kind, last_id)
-
-
-def _update_feed(hash_id):
-    print('Updating feed {}'.format(hash_id))
-    feed = _query_feed(hash_id)
-
-    page_size = int(feed.get('PageSize', DEFAULT_PAGE_SIZE))
-    last_id = feed.get('LastID', None)
-    episodes = feed.get('Episodes', [])
-    item_url = feed['ItemURL']
-
-    # Rebuild episode list from scratch
-    if not last_id:
-        episodes = []
-
-    start = time.time()
-    _, items, new_last_id = _get_updates(1, page_size, item_url, _get_format(feed), last_id)
-    end = time.time()
-
-    print('Got feed update: new {}, current {}. Update took: {}'.format(len(items), len(episodes), end-start))
-
-    # Update feed and submit back to Dynamo
-
-    unix_time = int(datetime.utcnow().timestamp())
-    feed['UpdatedAt'] = unix_time
-
-    if len(items) > 0:
-        episodes = items + episodes  # Prepand new episodes
-        del episodes[page_size:]  # Truncate list
-        feed['Episodes'] = episodes
-
-        # Update last seen video ID
-        feed['LastID'] = new_last_id
-
-        _update_feed_episodes(hash_id, feed)
-    else:
-        # Update last access field only
-        _update_feed_updated_at(hash_id, unix_time)
-
-
-def _query_feed(hash_id):
-    response = feeds_table.get_item(
-        Key={'HashID': hash_id},
-        ProjectionExpression='#prov,#type,#size,#fmt,#quality,#level,#id,#last_id,#episodes,#updated_at,#item_url',
-        ExpressionAttributeNames={
-            '#prov': 'Provider',
-            '#type': 'LinkType',
-            '#size': 'PageSize',
-            '#fmt': 'Format',
-            '#quality': 'Quality',
-            '#level': 'FeatureLevel',
-            '#id': 'ItemID',
-            '#last_id': 'LastID',
-            '#episodes': 'Episodes',
-            '#updated_at': 'UpdatedAt',
-            '#item_url': 'ItemURL',
-        },
-    )
-
-    item = response['Item']
-    return item
-
-
-def _update_feed_episodes(hash_id, feed):
-    feeds_table.update_item(
-        Key={
-            'HashID': hash_id,
-        },
-        UpdateExpression='SET #updated_at = :updated_at, #episodes = :episodes, #last_id = :last_id',
-        ExpressionAttributeNames={
-            '#updated_at': 'UpdatedAt',
-            '#episodes': 'Episodes',
-            '#last_id': 'LastID',
-        },
-        ExpressionAttributeValues={
-            ':updated_at': feed['UpdatedAt'],
-            ':episodes': feed['Episodes'],
-            ':last_id': feed['LastID'],
-        },
-        ReturnValues='NONE',
-    )
-
-
-def _update_feed_updated_at(hash_id, updated_at):
-    feeds_table.update_item(
-        Key={
-            'HashID': hash_id,
-        },
-        UpdateExpression='SET #updated_at = :updated_at',
-        ExpressionAttributeNames={
-            '#updated_at': 'UpdatedAt',
-        },
-        ExpressionAttributeValues={
-            ':updated_at': updated_at,
-        },
-        ReturnValues='NONE',
-    )
-
-
-def _get_format(feed):
-    fmt = feed.get('Format', 'video')
-    quality = feed.get('Quality', 'high')
 
+def _get_format(fmt, quality):
     if fmt == 'video':
         # Video
         if quality == 'high':
@@ -151,6 +29,9 @@ def _get_updates(start, count, url, fmt, last_id=None):
 
     end = start + count - 1
 
+    if not url:
+        raise ValueError('Invalid resource URL %s' % url)
+
     opts = {
         'playliststart': start,
         'playlistend': end,

cmd/updater/updater_test.py

@@ -7,10 +7,10 @@ TEST_URL = 'https://www.youtube.com/user/CNN/videos'
 class TestUpdater(unittest.TestCase):
     def test_get_updates(self):
         kinds = [
-            updater._get_format({'Format': 'video', 'Quality': 'high'}),
-            updater._get_format({'Format': 'video', 'Quality': 'low'}),
-            updater._get_format({'Format': 'audio', 'Quality': 'high'}),
-            updater._get_format({'Format': 'audio', 'Quality': 'low'}),
+            updater._get_format('video', 'high'),
+            updater._get_format('video', 'low'),
+            updater._get_format('audio', 'high'),
+            updater._get_format('audio', 'low'),
         ]
         for kind in kinds:
             with self.subTest(kind):
@@ -34,18 +34,3 @@ class TestUpdater(unittest.TestCase):
         feed, items, last_id = updater._get_updates(1, 1, TEST_URL, 'worstaudio')
         self.assertEqual(len(items), 1)
         self.assertEqual(items[0]['ID'], last_id)
-
-    @unittest.skip('heavy test, run manually')
-    def test_get_50(self):
-        _, items, last_id = updater.handler({
-            'url': 'https://www.youtube.com/channel/UCd6MoB9NC6uYN2grvUNT-Zg',
-            'start': 1,
-            'count': 50,
-            'kind': 'best[ext=mp4]',
-        }, None)
-        self.assertEqual(len(items), 50)
-        self.assertEqual(items[0]['ID'], last_id)
-
-    @unittest.skip
-    def test_update_feed(self):
-        updater._update_feed('86qZ')