From e3a6aa8addb999c4b4a72855835bb9461ce5acc0 Mon Sep 17 00:00:00 2001
From: Phil Jay <phil.jay@reece.com.au>
Date: Mon, 25 Nov 2019 12:18:45 +1100
Subject: [PATCH] Make localstack image name configurable

This allows folks to substitute their own localstack image, or point to
a local cache of Docker images - such as an Artifactory proxy/cache.
---
 pytest_localstack/__init__.py         | 4 ++++
 pytest_localstack/contrib/botocore.py | 4 ++++
 pytest_localstack/session.py          | 7 +++++--
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/pytest_localstack/__init__.py b/pytest_localstack/__init__.py
index 93eb873..a4aff6e 100644
--- a/pytest_localstack/__init__.py
+++ b/pytest_localstack/__init__.py
@@ -48,6 +48,7 @@ def session_fixture(
     kinesis_error_probability=0.0,
     dynamodb_error_probability=0.0,
     container_log_level=logging.DEBUG,
+    localstack_image="localstack/localstack",
     localstack_version="latest",
     auto_remove=True,
     pull_image=True,
@@ -95,6 +96,8 @@ def session_fixture(
             DynamoDB API responses.
         container_log_level (int, optional): The logging level to use
             for Localstack container logs. Defaults to :data:`logging.DEBUG`.
+        localstack_image (str, optional): The Docker image of Localstack
+            to use. Defaults to :const:`"localstack/localstack"`.
         localstack_version (str, optional): The version of the Localstack
             image to use. Defaults to :const:`"latest"`.
         auto_remove (bool, optional): If :obj:`True`, delete the Localstack
@@ -122,6 +125,7 @@ def _fixture(pytestconfig):
             kinesis_error_probability=kinesis_error_probability,
             dynamodb_error_probability=dynamodb_error_probability,
             container_log_level=container_log_level,
+            localstack_image=localstack_image,
             localstack_version=localstack_version,
             auto_remove=auto_remove,
             pull_image=pull_image,
diff --git a/pytest_localstack/contrib/botocore.py b/pytest_localstack/contrib/botocore.py
index e9424b0..a2b5533 100644
--- a/pytest_localstack/contrib/botocore.py
+++ b/pytest_localstack/contrib/botocore.py
@@ -324,6 +324,7 @@ def patch_fixture(
     kinesis_error_probability=0.0,
     dynamodb_error_probability=0.0,
     container_log_level=logging.DEBUG,
+    localstack_image="localstack/localstack",
     localstack_version="latest",
     auto_remove=True,
     pull_image=True,
@@ -370,6 +371,8 @@ def patch_fixture(
             DynamoDB API responses.
         container_log_level (int, optional): The logging level to use
             for Localstack container logs. Defaults to :data:`logging.DEBUG`.
+        localstack_image (str, optional): The Docker image of Localstack
+            to use. Defaults to :const:`"localstack/localstack"`.
         localstack_version (str, optional): The version of the Localstack
             image to use. Defaults to :const:`"latest"`.
         auto_remove (bool, optional): If :obj:`True`, delete the Localstack
@@ -397,6 +400,7 @@ def _fixture(pytestconfig):
             kinesis_error_probability=kinesis_error_probability,
             dynamodb_error_probability=dynamodb_error_probability,
             container_log_level=container_log_level,
+            localstack_image=localstack_image,
             localstack_version=localstack_version,
             auto_remove=auto_remove,
             pull_image=pull_image,
diff --git a/pytest_localstack/session.py b/pytest_localstack/session.py
index c96c007..48aa11c 100644
--- a/pytest_localstack/session.py
+++ b/pytest_localstack/session.py
@@ -203,6 +203,8 @@ class LocalstackSession(RunningSession):
             DynamoDB API responses.
         container_log_level (int, optional): The logging level to use
             for Localstack container logs. Defaults to :attr:`logging.DEBUG`.
+        localstack_image (str, optional): The Docker image of Localstack
+            to use. Defaults to `localstack/localstack`.
         localstack_version (str, optional): The version of the Localstack
             image to use. Defaults to `latest`.
         auto_remove (bool, optional): If True, delete the Localstack
@@ -216,7 +218,6 @@ class LocalstackSession(RunningSession):
 
     """
 
-    image_name = "localstack/localstack"
     factories = []
 
     def __init__(
@@ -227,6 +228,7 @@ def __init__(
         kinesis_error_probability=0.0,
         dynamodb_error_probability=0.0,
         container_log_level=logging.DEBUG,
+        localstack_image="localstack/localstack",
         localstack_version="latest",
         auto_remove=True,
         pull_image=True,
@@ -253,6 +255,7 @@ def __init__(
         )
 
         self.container_log_level = container_log_level
+        self.localstack_image = localstack_image
         self.localstack_version = localstack_version
         self.container_name = container_name or generate_container_name()
 
@@ -277,7 +280,7 @@ def start(self, timeout=60):
         logger.debug("%r running starting hooks", self)
         plugin.manager.hook.session_starting(session=self)
 
-        image_name = self.image_name + ":" + self.localstack_version
+        image_name = self.localstack_image + ":" + self.localstack_version
         if self.pull_image:
             logger.debug("Pulling docker image %r", image_name)
             self.docker_client.images.pull(image_name)