Refactor IntegrationClient into a separate file

Add AdminClient

Add Ephemeral client. Test for ephemeral client

Author: Val Brodsky

Makefile

@@ -64,6 +64,6 @@ test-ephemeral: build-image
 	docker run -it --rm -v ${PWD}:/usr/src -w /usr/src \
 		-e LABELBOX_TEST_ENVIRON="ephemeral" \
 		-e DA_GCP_LABELBOX_API_KEY=${DA_GCP_LABELBOX_API_KEY} \
-		-e LABELBOX_TEST_API_KEY_EPHEMERAL=${LABELBOX_TEST_API_KEY_LOCAL} \
-		-e FIXTURE_PROFILE=true \
+		-e SERVICE_API_KEY=${SERVICE_API_KEY} \
+		-e LABELBOX_TEST_BASE_URL="http://host.docker.internal:8080" \
 		local/labelbox-python:test pytest $(PATH_TO_TEST)

tests/integration/conftest.py

@@ -2,49 +2,34 @@
 from itertools import islice
 import json
 import os
-import re
 import sys
 import time
 import uuid
-from enum import Enum
 from types import SimpleNamespace
-from typing import Type, List, Union
+from typing import Type, List
 
 import pytest
 import requests
 
-from labelbox import Client, Dataset
+from labelbox import Dataset
 from labelbox import LabelingFrontend
 from labelbox import OntologyBuilder, Tool, Option, Classification, MediaType
 from labelbox.orm import query
 from labelbox.pagination import PaginatedCollection
 from labelbox.schema.annotation_import import LabelImport
 from labelbox.schema.enums import AnnotationImportState
 from labelbox.schema.invite import Invite
-from labelbox.schema.project import Project
 from labelbox.schema.quality_mode import QualityMode
 from labelbox.schema.queue_mode import QueueMode
 from labelbox.schema.user import User
+from support.integration_client import Environ, IntegrationClient, EphemeralClient, AdminClient
 
 IMG_URL = "https://picsum.photos/200/300.jpg"
 SMALL_DATASET_URL = "https://storage.googleapis.com/lb-artifacts-testing-public/sdk_integration_test/potato.jpeg"
 DATA_ROW_PROCESSING_WAIT_TIMEOUT_SECONDS = 30
 DATA_ROW_PROCESSING_WAIT_SLEEP_INTERNAL_SECONDS = 3
 
 
-class Environ(Enum):
-    LOCAL = 'local'
-    PROD = 'prod'
-    STAGING = 'staging'
-    ONPREM = 'onprem'
-    CUSTOM = 'custom'
-    STAGING_EU = 'staging-eu'
-    EPHEMERAL = 'ephemeral'  # Used for testing PRs with ephemeral environments
-
-
-EPHEMERAL_BASE_URL = "http://lb-api-public"
-
-
 @pytest.fixture(scope="session")
 def environ() -> Environ:
     """
@@ -60,69 +45,6 @@ def environ() -> Environ:
         raise Exception(f'Missing LABELBOX_TEST_ENVIRON in: {os.environ}')
 
 
-def graphql_url(environ: str) -> str:
-    if environ == Environ.PROD:
-        return 'https://api.labelbox.com/graphql'
-    elif environ == Environ.STAGING:
-        return 'https://api.lb-stage.xyz/graphql'
-    elif environ == Environ.STAGING_EU:
-        return 'https://api.eu-de.lb-stage.xyz/graphql'
-    elif environ == Environ.ONPREM:
-        hostname = os.environ.get('LABELBOX_TEST_ONPREM_HOSTNAME', None)
-        if hostname is None:
-            raise Exception(f"Missing LABELBOX_TEST_ONPREM_INSTANCE")
-        return f"{hostname}/api/_gql"
-    elif environ == Environ.CUSTOM:
-        graphql_api_endpoint = os.environ.get(
-            'LABELBOX_TEST_GRAPHQL_API_ENDPOINT')
-        if graphql_api_endpoint is None:
-            raise Exception(f"Missing LABELBOX_TEST_GRAPHQL_API_ENDPOINT")
-        return graphql_api_endpoint
-    elif environ == Environ.EPHEMERAL:
-        return f"{EPHEMERAL_BASE_URL}/graphql"
-    return 'http://host.docker.internal:8080/graphql'
-
-
-def rest_url(environ: str) -> str:
-    if environ == Environ.PROD:
-        return 'https://api.labelbox.com/api/v1'
-    elif environ == Environ.STAGING:
-        return 'https://api.lb-stage.xyz/api/v1'
-    elif environ == Environ.STAGING_EU:
-        return 'https://api.eu-de.lb-stage.xyz/api/v1'
-    elif environ == Environ.CUSTOM:
-        rest_api_endpoint = os.environ.get('LABELBOX_TEST_REST_API_ENDPOINT')
-        if rest_api_endpoint is None:
-            raise Exception(f"Missing LABELBOX_TEST_REST_API_ENDPOINT")
-        return rest_api_endpoint
-    elif environ == Environ.EPHEMERAL:
-        return f"{EPHEMERAL_BASE_URL}/api/v1"
-    return 'http://host.docker.internal:8080/api/v1'
-
-
-def admin_url(environ: str) -> Union[str, None]:
-    if environ == Environ.EPHEMERAL:
-        return f"{EPHEMERAL_BASE_URL}/admin/v1"
-
-    return 'http://host.docker.internal:8080/admin/v1'
-
-
-def testing_api_key(environ: str) -> str:
-    if environ == Environ.PROD:
-        return os.environ["LABELBOX_TEST_API_KEY_PROD"]
-    elif environ == Environ.STAGING:
-        return os.environ["LABELBOX_TEST_API_KEY_STAGING"]
-    elif environ == Environ.STAGING_EU:
-        return os.environ["LABELBOX_TEST_API_KEY_STAGING_EU"]
-    elif environ == Environ.ONPREM:
-        return os.environ["LABELBOX_TEST_API_KEY_ONPREM"]
-    elif environ == Environ.CUSTOM:
-        return os.environ["LABELBOX_TEST_API_KEY_CUSTOM"]
-    elif environ == Environ.EPHEMERAL:
-        return os.environ["LABELBOX_TEST_API_KEY_EPHEMERAL"]
-    return os.environ["LABELBOX_TEST_API_KEY_LOCAL"]
-
-
 def cancel_invite(client, invite_id):
     """
     Do not use. Only for testing.
@@ -172,121 +94,15 @@ def queries():
                            get_invites=get_invites)
 
 
-class IntegrationClient(Client):
-
-    def __init__(self, environ: str) -> None:
-        api_url = graphql_url(environ)
-        api_key = testing_api_key(environ)
-        rest_endpoint = rest_url(environ)
-        self._admin_endpoint = admin_url(environ)
-
-        super().__init__(api_key,
-                         api_url,
-                         enable_experimental=True,
-                         rest_endpoint=rest_endpoint)
-        self.queries = []
-
-    def execute(self, query=None, params=None, check_naming=True, **kwargs):
-        if check_naming and query is not None:
-            assert re.match(r"(?:query|mutation) \w+PyApi", query) is not None
-        self.queries.append((query, params))
-        return super().execute(query, params, **kwargs)
-
-    def create_organization(self) -> str:
-        endpoint = f"{self._admin_endpoint}/organizations/"
-        response = requests.post(
-            endpoint,
-            headers=self.headers,
-            json={"name": f"Test Org {uuid.uuid4()}"},
-        )
-        if response.status_code != requests.codes.created:
-            raise Exception("Failed to create ephemeral org, message: " +
-                            str(response.json()['message']))
-
-        return response.json()['id']
-
-    def create_user(self, organization_id) -> tuple[str, str]:
-        endpoint = f"{self._admin_endpoint}/user-identities/"
-        identity_id = f"e2e+{uuid.uuid4()}"
-
-        response = requests.post(
-            endpoint,
-            headers=self.headers,
-            json={
-                "identityId": identity_id,
-                "email": "email@email.com",
-                "name": f"tester{uuid.uuid4()}",
-                "verificationStatus": "VERIFIED",
-            },
-        )
-        if response.status_code != requests.codes.created:
-            raise Exception("Failed to create ephemeral org, message: " +
-                            str(response.json()['message']))
-
-        user_identity_id = response.json()['identityId']
-
-        endpoint = f"{self._admin_endpoint}/organizations/{organization_id}/users/"
-        response = requests.post(
-            endpoint,
-            headers=self.headers,
-            json={
-                "identityId": user_identity_id,
-                "organizationRole": "Admin"
-            },
-        )
-        if response.status_code != requests.codes.created:
-            raise Exception("Failed to create ephemeral org, message: " +
-                            str(response.json()['message']))
-
-        user_id = response.json()['id']
-
-        endpoint = f"{self._admin_endpoint}/users/{user_id}/token"
-        response = requests.get(
-            endpoint,
-            headers=self.headers,
-        )
-        if response.status_code != requests.codes.created:
-            raise Exception("Failed to create ephemeral org, message: " +
-                            str(response.json()['message']))
-
-        token = response["token"]
-
-        return user_id, token
-
-    def create_api_key_for_user(self, user_token) -> str:
-        key_name = f"test-key+{uuid.uuid4()}"
-        query = """
-            mutation CreateApiKey($name: String!) {
-                createApiKey(data: {name: $name}) {
-                    id
-                    jwt
-                    __typename
-                }
-            }
-        """
-        params = {"name": key_name}
-        req = self._make_gql_request(query=query, params=params)
-
-        return req["createApiKey"]["jwt"]
-        mutation_name = "deleteDataRowsByQuery"
-        query = """mutation DeleteDataRowsByQueryPyApi($searchQueryInput: SearchServiceQueryInput!)  {
-                        %s(where: {searchQuery: $searchQueryInput})
-                            { taskId }
-                        }
-                    """ % (mutation_name)
-        query_params = {
-            "searchQueryInput": {
-                "query": search_query,
-                "scope": None
-            }
-        }
-
-        res = self.execute(query, query_params, error_log_key="errors")
-        res = res[mutation_name]
+@pytest.fixture(scope="session")
+def admin_client(environ: str):
+    return AdminClient(environ)
 
 
 @pytest.fixture(scope="session")
 def client(environ: str):
+    if environ == Environ.EPHEMERAL:
+        return EphemeralClient()
     return IntegrationClient(environ)