Refactor _get_implict_client. Add ability to load custom clients from get_client. Refactor env variable to read auth headers as dict.

Author: fkdosilovic

.github/workflows/test.yml

@@ -32,10 +32,10 @@ jobs:
           JUDGE0_ATD_API_KEY: ${{ secrets.JUDGE0_ATD_API_KEY }}
           JUDGE0_RAPID_API_KEY: ${{ secrets.JUDGE0_RAPID_API_KEY }}
           JUDGE0_SULU_API_KEY: ${{ secrets.JUDGE0_SULU_API_KEY }}
-          JUDGE0_TEST_API_KEY: ${{ secrets.JUDGE0_TEST_API_KEY }}
-          JUDGE0_TEST_API_KEY_HEADER: ${{ secrets.JUDGE0_TEST_API_KEY_HEADER }}
-          JUDGE0_TEST_CE_ENDPOINT: ${{ secrets.JUDGE0_TEST_CE_ENDPOINT }}
-          JUDGE0_TEST_EXTRA_CE_ENDPOINT: ${{ secrets.JUDGE0_TEST_EXTRA_CE_ENDPOINT }}
+          JUDGE0_CE_AUTH_HEADERS: ${{ secrets.JUDGE0_CE_AUTH_HEADERS }}
+          JUDGE0_EXTRA_CE_AUTH_HEADERS: ${{ secrets.JUDGE0_EXTRA_CE_AUTH_HEADERS }}
+          JUDGE0_CE_ENDPOINT: ${{ secrets.JUDGE0_CE_ENDPOINT }}
+          JUDGE0_EXTRA_CE_ENDPOINT: ${{ secrets.JUDGE0_EXTRA_CE_ENDPOINT }}
         run: |
           source venv/bin/activate
           pytest -vv tests/

src/judge0/__init__.py

@@ -1,5 +1,7 @@
 import os
 
+from typing import Union
+
 from .api import (
     async_execute,
     async_run,
@@ -73,36 +75,25 @@ def _get_implicit_client(flavor: Flavor) -> Client:
     if flavor == Flavor.EXTRA_CE and JUDGE0_IMPLICIT_EXTRA_CE_CLIENT is not None:
         return JUDGE0_IMPLICIT_EXTRA_CE_CLIENT
 
-    from .clients import CE, EXTRA_CE
-
     try:
         from dotenv import load_dotenv
 
         load_dotenv()
     except:  # noqa: E722
         pass
 
-    if flavor == Flavor.CE:
-        client_classes = CE
-    else:
-        client_classes = EXTRA_CE
+    # Let's check if we can find a self-hosted client.
+    client = _get_custom_client(flavor)
 
     # Try to find one of the predefined keys JUDGE0_{SULU,RAPID,ATD}_API_KEY
     # in environment variables.
-    client = None
-    for client_class in client_classes:
-        api_key = os.getenv(client_class.API_KEY_ENV)
-        if api_key is not None:
-            client = client_class(api_key)
-            break
+    if client is None:
+        client = _get_predefined_client(flavor)
 
     # If we didn't find any of the possible predefined keys, initialize
     # the preview Sulu client based on the flavor.
     if client is None:
-        if flavor == Flavor.CE:
-            client = SuluJudge0CE(retry_strategy=RegularPeriodRetry(0.5))
-        else:
-            client = SuluJudge0ExtraCE(retry_strategy=RegularPeriodRetry(0.5))
+        client = _get_preview_client(flavor)
 
     if flavor == Flavor.CE:
         JUDGE0_IMPLICIT_CE_CLIENT = client
@@ -112,6 +103,57 @@ def _get_implicit_client(flavor: Flavor) -> Client:
     return client
 
 
+def _get_preview_client(flavor: Flavor) -> Union[SuluJudge0CE, SuluJudge0ExtraCE]:
+    if flavor == Flavor.CE:
+        return SuluJudge0CE(retry_strategy=RegularPeriodRetry(0.5))
+    else:
+        return SuluJudge0ExtraCE(retry_strategy=RegularPeriodRetry(0.5))
+
+
+def _get_custom_client(flavor: Flavor) -> Union[Client, None]:
+    ce_endpoint = os.getenv("JUDGE0_CE_ENDPOINT")
+    ce_auth_header = os.getenv("JUDGE0_CE_AUTH_HEADERS")
+    extra_ce_endpoint = os.getenv("JUDGE0_EXTRA_CE_ENDPOINT")
+    extra_ce_auth_header = os.getenv("JUDGE0_EXTRA_CE_AUTH_HEADERS")
+
+    if flavor == Flavor.CE and ce_endpoint is not None and ce_auth_header is not None:
+        return Client(
+            endpoint=ce_endpoint,
+            auth_headers=eval(ce_auth_header),
+        )
+
+    if (
+        flavor == Flavor.EXTRA_CE
+        and extra_ce_endpoint is not None
+        and extra_ce_auth_header is not None
+    ):
+        return Client(
+            endpoint=extra_ce_endpoint,
+            auth_headers=eval(extra_ce_auth_header),
+        )
+
+    return None
+
+
+def _get_predefined_client(flavor: Flavor) -> Union[Client, None]:
+    from .clients import CE, EXTRA_CE
+
+    if flavor == Flavor.CE:
+        client_classes = CE
+    else:
+        client_classes = EXTRA_CE
+
+    for client_class in client_classes:
+        api_key = os.getenv(client_class.API_KEY_ENV)
+        if api_key is not None:
+            client = client_class(api_key)
+            break
+    else:
+        client = None
+
+    return client
+
+
 CE = Flavor.CE
 EXTRA_CE = Flavor.EXTRA_CE
 

tests/conftest.py

@@ -10,34 +10,24 @@
 
 @pytest.fixture(scope="session")
 def judge0_ce_client():
-    api_key = os.getenv("JUDGE0_TEST_API_KEY")
-    api_key_header = os.getenv("JUDGE0_TEST_API_KEY_HEADER")
-    endpoint = os.getenv("JUDGE0_TEST_CE_ENDPOINT")
+    endpoint = os.getenv("JUDGE0_CE_ENDPOINT")
+    auth_headers = os.getenv("JUDGE0_CE_AUTH_HEADERS")
 
-    if api_key is None or api_key_header is None or endpoint is None:
+    if endpoint is None or auth_headers is None:
         return None
     else:
-        client = clients.Client(
-            endpoint=endpoint,
-            auth_headers={api_key_header: api_key},
-        )
-        return client
+        return clients.Client(endpoint=endpoint, auth_headers=eval(auth_headers))
 
 
 @pytest.fixture(scope="session")
 def judge0_extra_ce_client():
-    api_key = os.getenv("JUDGE0_TEST_API_KEY")
-    api_key_header = os.getenv("JUDGE0_TEST_API_KEY_HEADER")
-    endpoint = os.getenv("JUDGE0_TEST_EXTRA_CE_ENDPOINT")
+    endpoint = os.getenv("JUDGE0_EXTRA_CE_ENDPOINT")
+    auth_headers = os.getenv("JUDGE0_EXTRA_CE_AUTH_HEADERS")
 
-    if api_key is None or api_key_header is None or endpoint is None:
+    if endpoint is None or auth_headers is None:
         return None
     else:
-        client = clients.Client(
-            endpoint=endpoint,
-            auth_headers={api_key_header: api_key},
-        )
-        return client
+        return clients.Client(endpoint=endpoint, auth_headers=eval(auth_headers))
 
 
 @pytest.fixture(scope="session")