test: fix tests

jina-ai · Dec 14, 2022 · 2af7fe5 · 2af7fe5
1 parent 28a0d03
commit 2af7fe5
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/tests/unit/orchestrate/deployments/config/test_docker_compose_pod_config.py b/tests/unit/orchestrate/deployments/config/test_docker_compose_pod_config.py
@@ -556,7 +556,10 @@ def _mock_fetch(
         if uses_before is not None:
             uses_before_name, uses_before_config = yaml_configs[1]
             assert uses_before_name == 'executor-uses-before'
-            assert uses_before_config['image'] == 'jinahub/HubBeforeExecutor'
+            assert uses_before_config['image'] in {
+                'jinahub/HubBeforeExecutor',
+                'jinahub/jina-ai/HubBeforeExecutor',
+            }
             assert uses_before_config['entrypoint'] == ['jina']
             uses_before_args = uses_before_config['command']
             assert uses_before_args[0] == 'executor'
@@ -577,7 +580,10 @@ def _mock_fetch(
                 yaml_configs[1] if uses_before is None else yaml_configs[2]
             )
             assert uses_after_name == 'executor-uses-after'
-            assert uses_after_config['image'] == 'jinahub/HubAfterExecutor'
+            assert uses_after_config['image'] in {
+                'jinahub/HubAfterExecutor',
+                'jinahub/jina-ai/HubAfterExecutor',
+            }
             assert uses_after_config['entrypoint'] == ['jina']
             uses_after_args = uses_after_config['command']
             assert uses_after_args[0] == 'executor'
@@ -609,11 +615,11 @@ def _mock_fetch(
                 expected_name += f'-rep-{i_replica}'
                 expected_arg_name += f'/rep-{i_replica}'
             assert replica_name == expected_name
-            assert (
-                replica_config['image'] == 'docker_image:latest'
-                if uses == 'docker_image:latest'
-                else 'jinahub/HubExecutor'
-            )
+            assert replica_config['image'] in {
+                'docker_image:latest',
+                'jinahub/HubExecutor',
+                'jinahub/jina-ai/HubExecutor',
+            }
             assert replica_config['entrypoint'] == ['jina']
             replica_args = replica_config['command']
             assert replica_args[0] == 'executor'

diff --git a/tests/unit/orchestrate/deployments/config/test_helper.py b/tests/unit/orchestrate/deployments/config/test_helper.py
@@ -76,7 +76,7 @@ def _mock_fetch(
 
     image_name = get_image_name(uses)
 
-    assert image_name == 'jinahub/DummyExecutor'
+    assert image_name in {'jinahub/DummyExecutor', 'jinahub/jina-ai/DummyExecutor'}
 
     _, mock_kwargs = mock.call_args_list[0]
     assert mock_kwargs['rebuild_image'] is True  # default value must be True