From 1b4fc8667714ef4ce9f326bd14f795fc2417ecb9 Mon Sep 17 00:00:00 2001
From: Terry Truong <terry06890@gmail.com>
Date: Sat, 1 Oct 2022 23:14:08 +1000
Subject: Add per-event-category image limit

---
 backend/tests/enwiki/test_gen_img_data.py | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

(limited to 'backend/tests/enwiki/test_gen_img_data.py')

diff --git a/backend/tests/enwiki/test_gen_img_data.py b/backend/tests/enwiki/test_gen_img_data.py
index 019b757..dc77fe6 100644
--- a/backend/tests/enwiki/test_gen_img_data.py
+++ b/backend/tests/enwiki/test_gen_img_data.py
@@ -19,6 +19,22 @@ class TestGetInputPageIdsFromDb(unittest.TestCase):
 				{
 					(1, 'Belgium', 2389729, None, None, None, 2, 'country'),
 					(2, 'George Washington', 2353711, None, 2378478, None, 2, 'human'),
+					(3, 'Douglas Adams', 2434082, None, 2452040, None, 2, 'human'),
+					(4, 'World War II', 2429507, None, 2431700, None, 2, 'event'),
+					(5, 'Marie Curie', 2403277, None, 2427622, None, 2, 'human'),
+				}
+			)
+			# Create temp pageviews db
+			pageviewDb = os.path.join(tempDir, 'pageview_data.db')
+			createTestDbTable(
+				pageviewDb,
+				'CREATE TABLE views (title TEXT PRIMARY KEY, id INT, views INT)',
+				'INSERT INTO views VALUES (?, ?, ?)',
+				{
+					('George Washington', 2, 8),
+					('Marie Curie', 5, 10),
+					('Douglas Adams', 3, 5),
+					('Belgium', 1, 100),
 				}
 			)
 			# Create temp dump-index db
@@ -30,13 +46,15 @@ class TestGetInputPageIdsFromDb(unittest.TestCase):
 				{
 					('Belgium',10,0,-1),
 					('George Washington',20,0,-1),
+					('Douglas Adamns',30,0,-1),
+					('Marie Curie',50,0,-1),
 					('Autism',25,0,-1),
 				}
 			)
 			# Run
-			pageIds = getInputPageIdsFromDb(dbFile, indexDb)
+			pageIds = getInputPageIdsFromDb(dbFile, pageviewDb, indexDb, 2)
 			# Check
-			self.assertEqual(pageIds, {10, 20})
+			self.assertEqual(pageIds, {50, 20, 10})
 
 class TestGenData(unittest.TestCase):
 	def test_gen(self):
-- 
cgit v1.2.3