| 1 | import unittest |
| 2 | |
| 3 | from lib import entity_extract |
| 4 | |
| 5 | |
| 6 | class TestExtractSubreddits(unittest.TestCase): |
| 7 | def test_extracts_primary_subreddit(self): |
| 8 | items = [{"subreddit": "r/MachineLearning"}] |
| 9 | result = entity_extract._extract_subreddits(items) |
| 10 | self.assertIn("MachineLearning", result) |
| 11 | |
| 12 | def test_extracts_subreddit_without_prefix(self): |
| 13 | items = [{"subreddit": "localLLaMA"}] |
| 14 | result = entity_extract._extract_subreddits(items) |
| 15 | self.assertIn("localLLaMA", result) |
| 16 | |
| 17 | def test_extracts_cross_references_from_comment_insights(self): |
| 18 | items = [ |
| 19 | { |
| 20 | "subreddit": "technology", |
| 21 | "comment_insights": ["check out r/MachineLearning and r/LocalLLaMA for more"], |
| 22 | } |
| 23 | ] |
| 24 | result = entity_extract._extract_subreddits(items) |
| 25 | self.assertIn("MachineLearning", result) |
| 26 | self.assertIn("LocalLLaMA", result) |
| 27 | |
| 28 | def test_extracts_cross_references_from_top_comments(self): |
| 29 | items = [ |
| 30 | { |
| 31 | "subreddit": "AI", |
| 32 | "top_comments": [ |
| 33 | {"excerpt": "see r/StableDiffusion for image gen stuff"}, |
| 34 | ], |
| 35 | } |
| 36 | ] |
| 37 | result = entity_extract._extract_subreddits(items) |
| 38 | self.assertIn("StableDiffusion", result) |
| 39 | |
| 40 | def test_ranks_by_frequency(self): |
| 41 | items = [ |
| 42 | {"subreddit": "A"}, |
| 43 | {"subreddit": "A"}, |
| 44 | {"subreddit": "B"}, |
| 45 | ] |
| 46 | result = entity_extract._extract_subreddits(items) |
| 47 | self.assertEqual(result[0], "A") |
| 48 | |
| 49 | def test_empty_items(self): |
| 50 | self.assertEqual([], entity_extract._extract_subreddits([])) |
| 51 | |
| 52 | if __name__ == "__main__": |
| 53 | unittest.main() |
| 54 |