unit test for token usage

PeeachPie · PeeachPie · commit 42926f045ca7 · 2025-03-29T15:22:16.000+03:00
diff --git a/git_logger.py b/git_logger.py
@@ -2,8 +2,6 @@
 from time import sleep
 
 TIMEDELTA = 0.05
-TIMEZONE = 'Europe/Moscow'
-
 
 def login(source, token, base_url):
     client = RepositoryFactory.create_api(source, token, base_url)
diff --git a/interface_wrapper.py b/interface_wrapper.py
@@ -3,7 +3,7 @@
 from dataclasses import dataclass
 import logging
 
-from github import Github
+from github import Github, Auth
 from pyforgejo import PyforgejoApi
 
 # Настройка логирования
@@ -182,7 +182,7 @@ def create_api(
         from ForgejoRepoAPI import ForgejoRepoAPI
 
         if source == 'github':
-            return GitHubRepoAPI(Github(token))
+            return GitHubRepoAPI(Github(auth=Auth.Token(token)))
         elif source == 'forgejo':
             if not isinstance(base_url, str):
                 raise ValueError(
diff --git a/main.py b/main.py
@@ -118,6 +118,42 @@ def parse_args():
     return parser.parse_args()
 
 
+def run(args, binded_repos, repos_for_wiki=None):
+    start = parse_time(args.start.split('-'))
+    finish = parse_time(args.finish.split('-'))
+
+    if args.commits:
+        commits_parser.log_commits(
+            binded_repos, args.out, start, finish, args.branch, args.forks_include
+        )
+    if args.pull_requests:
+        pull_requests_parser.log_pull_requests(
+            binded_repos,
+            args.out,
+            start,
+            finish,
+            args.forks_include,
+            args.pr_comments,
+        )
+    if args.issues:
+        issues_parser.log_issues(
+            binded_repos, args.out, start, finish, args.forks_include
+        )
+    if args.invites:
+        invites_parser.log_invitations(
+            binded_repos,
+            args.out,
+        )
+    if args.contributors:
+        contributors_parser.log_contributors(binded_repos, args.out, args.forks_include)
+    if args.wikis:
+        wiki_parser.wiki_parser(repos_for_wiki, args.download_repos, args.out)
+    if args.export_google_sheets:
+        export_sheets.write_data_to_table(
+            args.out, args.google_token, args.table_id, args.sheet_id
+        )
+
+
 def main():
     args = parse_args()
 
@@ -128,51 +164,16 @@ def main():
 
     repositories = git_logger.get_repos_from_file(args.list)
 
-    csv_name = args.out
-    path_drepo = args.download_repos
-    fork_flag = args.forks_include
-    log_pr_comments = args.pr_comments
-
     print(repositories)
 
     try:
         clients = git_logger.Clients("github", tokens)
         binded_repos = git_logger.get_next_binded_repo(clients, repositories)
-        start = parse_time(args.start.split('-'))
-        finish = parse_time(args.finish.split('-'))
     except Exception as e:
         print(e)
         print(traceback.print_exc())
     else:
-        if args.commits:
-            commits_parser.log_commits(
-                binded_repos, csv_name, start, finish, args.branch, fork_flag
-            )
-        if args.pull_requests:
-            pull_requests_parser.log_pull_requests(
-                binded_repos,
-                csv_name,
-                start,
-                finish,
-                fork_flag,
-                log_pr_comments,
-            )
-        if args.issues:
-            issues_parser.log_issues(binded_repos, csv_name, start, finish, fork_flag)
-        if args.invites:
-            invites_parser.log_invitations(
-                binded_repos,
-                csv_name,
-            )
-        if args.wikis:
-            wiki_parser.wiki_parser(repositories, path_drepo, csv_name)
-        if args.contributors:
-            contributors_parser.log_contributors(binded_repos, csv_name, fork_flag)
-        if args.export_google_sheets:
-            export_sheets.write_data_to_table(
-                csv_name, args.google_token, args.table_id, args.sheet_id
-            )
-
+        run(args, binded_repos)
 
 if __name__ == '__main__':
     main()
diff --git a/test_token_usage.py b/test_token_usage.py
@@ -1,27 +1,12 @@
 import unittest
 import argparse
 import sys
-from utils import parse_time
-from datetime import datetime
-from interface_wrapper import RepositoryFactory, IRepositoryAPI
-
-import git_logger
-
-from repo_parser import (
-    commits_parser,
-    contributors_parser,
-    pull_requests_parser,
-    invites_parser,
-    issues_parser,
-    wiki_parser,
-)
 
+from main import run
 
-def fix_rate_limit(clients: git_logger.Clients):
-    return [c['client'].get_rate_limiting() for c in clients.clients]
-
+import git_logger
 
-def parse_args():
+def parse_args(args):
     parser = argparse.ArgumentParser()
     parser.add_argument('--tt1', type=str, required=True, help='first test token')
     parser.add_argument('--tt2', type=str, required=True, help='second test token')
@@ -36,45 +21,117 @@ def parse_args():
 
     parser.add_argument('-o', '--out', type=str, required=True, help='output filename')
 
-    return parser.parse_args()
-
+    return parser.parse_args(args)
 
-class TestCommitsParser(unittest.TestCase):
+class TestTokenUsage(unittest.TestCase):
     def setUp(self):
-        args = parse_args()
-        print(args)
+        test_args = parse_args(sys.argv[1:])
+
+        self.tokens = [test_args.tt1, test_args.tt2]
+        self.repo = test_args.repo
+        self.output_csv = test_args.out
+
+        self.args = argparse.Namespace(
+            commits=False,
+            issues=False,
+            pull_requests=False,
+            wikis=False,
+            contributors=False,
+            invites=False,
+            start="2000/01/01-00:00:00",
+            finish="2400/01/01-00:00:00",
+            branch="default",
+            forks_include=False,
+            pr_comments=False,
+            export_google_sheets=False,
+            out=test_args.out,
+        )
+
+    @staticmethod
+    def _get_rate_limit(clients: git_logger.Clients):
+        return [c['client'].get_rate_limiting()[0] for c in clients.clients]
+
+    @staticmethod
+    def _is_only_one_token_used(limit_start, limit_finish):
+        return (bool(limit_start[0] - limit_finish[0])
+            != bool(limit_start[1] - limit_finish[1]))
+    
+    @staticmethod
+    def _is_max_token_used(limit_start, limit_finish):
+        if limit_start[0] - limit_finish[0]:
+            return limit_start[0] == max(limit_start)
+        else:
+            return limit_start[1] == max(limit_start)
+    
+    @staticmethod
+    def _change_tokens_order(tokens, key):
+        key %= len(tokens)
+        return tokens[key:] + tokens[:key]
+
+    def _get_usage(self, binded_repos, clients):
+        limit_start = self._get_rate_limit(clients)
+
+        run(self.args, binded_repos)
 
-        self.token1 = args.tt1
-        self.token2 = args.tt2
-        self.repo = args.test_repo
-        self.output_csv = args.out
+        limit_finish = self._get_rate_limit(clients)
 
-        self.start = parse_time('2000/01/01-00:00:00')
-        self.finish = parse_time('2400/01/01-00:00:00')
-        self.branch = 'default'
-        self.fork_flag = False
+        return limit_start, limit_finish
 
     def test_commits_parser(self):
-        clients1 = git_logger.Clients("github", [self.token1, self.token2])
-        binded_repos1 = git_logger.get_next_binded_repo(clients1, [self.test_repo])
-
-        rate_limit_start = fix_rate_limit(clients1)
-
-        commits_parser.log_commits(
-            binded_repos1,
-            self.output_csv,
-            self.start,
-            self.finish,
-            self.branch,
-            self.fork_flag,
-        )
+        self.args.commits = True
+        self.assertTrue(False)
+        for i in range(2):
+            clients = git_logger.Clients("github", self._change_tokens_order(self.tokens, i))
+            binded_repos = git_logger.get_next_binded_repo(clients, [self.repo])
+
+            limit_start, limit_finish = self._get_usage(binded_repos, clients)
+
+            self.assertTrue(self._is_only_one_token_used(limit_start, limit_finish))
+            self.assertTrue(self._is_max_token_used(limit_start, limit_finish))
+
+    def test_contributors_parser(self):
+        self.args.contributors = True
+        for i in range(2):
+            clients = git_logger.Clients("github", self._change_tokens_order(self.tokens, i))
+            binded_repos = git_logger.get_next_binded_repo(clients, [self.repo])
+
+            limit_start, limit_finish = self._get_usage(binded_repos, clients)
+
+            self.assertTrue(self._is_only_one_token_used(limit_start, limit_finish))
+            self.assertTrue(self._is_max_token_used(limit_start, limit_finish))
+
+    def test_issues_parser(self):
+        self.args.issues = True
+        for i in range(2):
+            clients = git_logger.Clients("github", self._change_tokens_order(self.tokens, i))
+            binded_repos = git_logger.get_next_binded_repo(clients, [self.repo])
+
+            limit_start, limit_finish = self._get_usage(binded_repos, clients)
+
+            self.assertTrue(self._is_only_one_token_used(limit_start, limit_finish))
+            self.assertTrue(self._is_max_token_used(limit_start, limit_finish))
+
+    def test_invites_parser(self):
+        self.args.invites = True
+        for i in range(2):
+            clients = git_logger.Clients("github", self._change_tokens_order(self.tokens, i))
+            binded_repos = git_logger.get_next_binded_repo(clients, [self.repo])
+
+            limit_start, limit_finish = self._get_usage(binded_repos, clients)
 
-        rate_limit_finish = fix_rate_limit(clients1)
+            self.assertTrue(self._is_only_one_token_used(limit_start, limit_finish))
+            self.assertTrue(self._is_max_token_used(limit_start, limit_finish))
 
-        print(rate_limit_start, rate_limit_finish)
+    def test_pull_requests_parser(self):
+        self.args.pull_requests = True
+        for i in range(2):
+            clients = git_logger.Clients("github", self._change_tokens_order(self.tokens, i))
+            binded_repos = git_logger.get_next_binded_repo(clients, [self.repo])
 
-        pass
+            limit_start, limit_finish = self._get_usage(binded_repos, clients)
 
+            self.assertTrue(self._is_only_one_token_used(limit_start, limit_finish))
+            self.assertTrue(self._is_max_token_used(limit_start, limit_finish))
 
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main(argv=[sys.argv[0]])
diff --git a/utils.py b/utils.py
@@ -1,14 +1,13 @@
 import csv
 from datetime import datetime
 import pytz
-import git_logger
 
 TITLE_LEN = 80
 MIN_SIDE_PADDING = 4
 SIDE_WHITE_SPACES = 1
 
-
 class logger:
+    #TODO: отключение вывода в stdout
     @staticmethod
     def log_title(title: str, title_len: int = TITLE_LEN):
         final_len = max(
@@ -49,6 +48,7 @@ def log_error(error: str):
     def log_warning(warning: str):
         pass
 
+TIMEZONE = 'Europe/Moscow'
 
 def parse_time(datetime_str) -> datetime:
     start = (
@@ -65,4 +65,4 @@ def parse_time(datetime_str) -> datetime:
         minute=start[4],
         second=start[5],
     )
-    return start_datetime.astimezone(pytz.timezone(git_logger.TIMEZONE))
+    return start_datetime.astimezone(pytz.timezone(TIMEZONE))
