From 7eddc4a03af230bb40d0b5ef1c304c2828397fe7 Mon Sep 17 00:00:00 2001
From: Tzu-ping Chung <uranusjr@gmail.com>
Date: Sun, 20 Aug 2023 14:55:10 +0800
Subject: [PATCH] Avoid WSL2 ones when finding a context for Breeze

---
 .../utils/docker_command_utils.py             | 34 ++++++++++++-------
 dev/breeze/tests/test_docker_command_utils.py | 33 +++++++++++++++---
 2 files changed, 49 insertions(+), 18 deletions(-)

diff --git a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py
index 37a2ba3f8a9f5..d711d3d933814 100644
--- a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py
+++ b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py
@@ -17,6 +17,7 @@
 """Various utils to prepare docker and docker compose commands."""
 from __future__ import annotations
 
+import json
 import os
 import re
 import sys
@@ -800,23 +801,30 @@ def autodetect_docker_context():
 
     :return: name of the docker context to use
     """
-    output = run_command(["docker", "context", "ls", "-q"], capture_output=True, check=False, text=True)
-    if output.returncode != 0:
+    result = run_command(
+        ["docker", "context", "ls", "--format=json"],
+        capture_output=True,
+        check=False,
+        text=True,
+    )
+    if result.returncode != 0:
         get_console().print("[warning]Could not detect docker builder. Using default.[/]")
         return "default"
-    context_list = output.stdout.splitlines()
-    if not context_list:
+    known_contexts = {info["Name"]: info for info in json.loads(result.stdout)}
+    if not known_contexts:
         get_console().print("[warning]Could not detect docker builder. Using default.[/]")
         return "default"
-    elif len(context_list) == 1:
-        get_console().print(f"[info]Using {context_list[0]} as context.[/]")
-        return context_list[0]
-    else:
-        for preferred_context in PREFERRED_CONTEXTS:
-            if preferred_context in context_list:
-                get_console().print(f"[info]Using {preferred_context} as context.[/]")
-                return preferred_context
-    fallback_context = context_list[0]
+    for preferred_context_name in PREFERRED_CONTEXTS:
+        try:
+            context = known_contexts[preferred_context_name]
+        except KeyError:
+            continue
+        # On Windows, some contexts are used for WSL2. We don't want to use those.
+        if context["DockerEndpoint"] == "npipe:////./pipe/dockerDesktopLinuxEngine":
+            continue
+        get_console().print(f"[info]Using {preferred_context_name} as context.[/]")
+        return preferred_context_name
+    fallback_context = next(iter(known_contexts))
     get_console().print(
         f"[warning]Could not use any of the preferred docker contexts {PREFERRED_CONTEXTS}.\n"
         f"Using {fallback_context} as context.[/]"
diff --git a/dev/breeze/tests/test_docker_command_utils.py b/dev/breeze/tests/test_docker_command_utils.py
index 4b1513e30fc71..d16dada5903a5 100644
--- a/dev/breeze/tests/test_docker_command_utils.py
+++ b/dev/breeze/tests/test_docker_command_utils.py
@@ -195,19 +195,42 @@ def test_check_docker_compose_version_ok(mock_get_console, mock_run_command):
     )
 
 
+def _fake_ctx(name: str) -> dict[str, str]:
+    return {
+        "Name": name,
+        "DockerEndpoint": f"unix://{name}",
+    }
+
+
 @pytest.mark.parametrize(
     "context_output, selected_context, console_output",
     [
         (
+            json.dumps([_fake_ctx("default")]),
             "default",
+            "[info]Using default as context",
+        ),
+        ("[]", "default", "[warning]Could not detect docker builder"),
+        (
+            json.dumps([_fake_ctx("a"), _fake_ctx("b")]),
+            "a",
+            "[warning]Could not use any of the preferred docker contexts",
+        ),
+        (
+            json.dumps([_fake_ctx("a"), _fake_ctx("desktop-linux")]),
+            "desktop-linux",
+            "[info]Using desktop-linux as context",
+        ),
+        (
+            json.dumps([_fake_ctx("a"), _fake_ctx("default")]),
             "default",
             "[info]Using default as context",
         ),
-        ("", "default", "[warning]Could not detect docker builder"),
-        ("a\nb", "a", "[warning]Could not use any of the preferred docker contexts"),
-        ("a\ndesktop-linux", "desktop-linux", "[info]Using desktop-linux as context"),
-        ("a\ndefault", "default", "[info]Using default as context"),
-        ("a\ndefault\ndesktop-linux", "desktop-linux", "[info]Using desktop-linux as context"),
+        (
+            json.dumps([_fake_ctx("a"), _fake_ctx("default"), _fake_ctx("desktop-linux")]),
+            "desktop-linux",
+            "[info]Using desktop-linux as context",
+        ),
     ],
 )
 def test_autodetect_docker_context(context_output: str, selected_context: str, console_output: str):