[Python] Allow externally created builds in Python / Robot

visual-python/src/saucelabs_visual/client.py

@@ -11,16 +11,31 @@
 
 from saucelabs_visual.regions import Region
 from saucelabs_visual.typing import IgnoreRegion, FullPageConfig, DiffingMethod, BuildStatus, \
-    DiffingOptions, IgnoreElementRegion
+    DiffingOptions, IgnoreElementRegion, BuildMode
 
 PKG_VERSION = '0.0.12'
 
 
 class SauceLabsVisual:
     _client: Client = None
     build_id: Union[str, None] = None
+    """
+    The UUID of the current build managed by this client. 
+    """
+    is_build_external: bool = False
+    """
+    Whether this build was created externally and passed here. When 'True', we should not perform
+    any mutations on the build itself assuming the lifecycle is managed externally.
+    """
     build_url: Union[str, None] = None
+    """
+    The URL for the current build in the Sauce Labs UI.
+    """
     meta_cache: dict = {}
+    """
+    A cache holding session metadata for a job so we don't have to re-query for every snapshot for
+    the same session ID.
+    """
     region: Region = None
 
     @property
@@ -74,6 +89,15 @@ def create_build(
             before timing out and reporting an error in the UI.
         :return:
         """
+
+        custom_build = self._get_external_build()
+
+        if custom_build is not None:
+            self.build_id = custom_build['id']
+            self.build_url = custom_build['url']
+            self.is_build_external = True
+            return custom_build
+
         query = gql(
             # language=GraphQL
             """
@@ -94,6 +118,7 @@ def create_build(
                     keepAliveTimeout: $keepAliveTimeout,
                 }){
                     id
+                    mode
                     url
                 }
             }
@@ -112,7 +137,15 @@ def create_build(
         self.build_url = build['createBuild']['url']
         return build
 
-    def finish_build(self):
+    def finish_build(self, build_id: Union[str, None] = None):
+        """
+        Finish the build associated with this client or, optionally, one supplied with 'build_id'.
+        :param build_id: An optional param to allow finishing an external build.
+        :return:
+        """
+        if self.is_build_external:
+            return
+
         query = gql(
             # language=GraphQL
             """
@@ -125,7 +158,7 @@ def finish_build(self):
             }
             """
         )
-        values = {"id": self.build_id}
+        values = {"id": build_id or self.build_id}
         self.meta_cache.clear()
         return self.client.execute(query, variable_values=values)
 
@@ -261,6 +294,74 @@ def create_snapshot_from_webdriver(
         }
         return self.client.execute(query, variable_values=values)
 
+    def _get_external_build(self) -> Union[dict, None]:
+        """
+        Check if we've been provided environment variables to use an external build instead of one
+        handled by the lifecycle in this client. If provided, query for the build and set the client
+        up to use it.
+        :return:
+        """
+        env_build_id = environ.get('SAUCE_VISUAL_BUILD_ID')
+        env_custom_id = environ.get('SAUCE_VISUAL_CUSTOM_ID')
+        build: Union[dict, None] = None
+
+        def get_build_meta(result: dict, error_message: str):
+            result_build = result['result']
+            if result_build is None:
+                raise RuntimeError(
+                    f"Sauce Labs Visual: unable to fetch build for {error_message}. Build not found."
+                )
+
+            if result_build['mode'] == BuildMode.COMPLETED.value:
+                raise RuntimeError(
+                    "Sauce Labs Visual: cannot add more screenshots since the build is already "
+                    "completed"
+                )
+
+            return result_build
+
+        if env_build_id:
+            query = gql(
+                # language=GraphQL
+                """
+                query buildById($buildId: UUID!) {
+                    result: build(id: $buildId) {
+                        id
+                        mode
+                        url
+                    }
+                }
+                """
+            )
+            build = get_build_meta(
+                self.client.execute(query, variable_values={
+                    "buildId": env_build_id,
+                }),
+                "SAUCE_VISUAL_BUILD_ID"
+            )
+
+        if env_custom_id:
+            query = gql(
+                # language=GraphQL
+                """
+                query buildById($customId: String!) {
+                    result: buildByCustomId(customId: $customId) {
+                        id
+                        mode
+                        url
+                    }
+                }
+                """
+            )
+            build = get_build_meta(
+                self.client.execute(query, variable_values={
+                    "customId": env_custom_id,
+                }),
+                "SAUCE_VISUAL_CUSTOM_ID"
+            )
+
+        return build
+
     def get_build_status(
             self,
             wait: bool = True,