From 7ee6b89aee34401f8ea26f67df777bfd6ffd1ed4 Mon Sep 17 00:00:00 2001
From: Chris White <chris@prefect.io>
Date: Fri, 8 Nov 2024 07:07:47 -0800
Subject: [PATCH] Fix failing format tests

---
 test_oss_cloud_api_compatibility.py | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/test_oss_cloud_api_compatibility.py b/test_oss_cloud_api_compatibility.py
index 2ebf060..7ba2e10 100644
--- a/test_oss_cloud_api_compatibility.py
+++ b/test_oss_cloud_api_compatibility.py
@@ -5,9 +5,16 @@
 import pytest
 
 
+PREFECT_V2 = False
+
+
 @pytest.fixture
 def oss_schema():
-    return load_schema("oss_schema.json")
+    schema = load_schema("oss_schema.json")
+    if schema["info"]["version"].startswith("2"):
+        global PREFECT_V2
+        PREFECT_V2 = True
+    return schema
 
 
 @pytest.fixture
@@ -252,11 +259,21 @@ def extract_types(d):
             return {item.get("type") for item in d["anyOf"] if item.get("type")}
         return set()
 
+    def extract_format(d):
+        if "format" in d:
+            return d["format"]
+        # in practice, this will have only one format
+        elif "anyOf" in d:
+            for option in d["anyOf"]:
+                if option.get("format"):
+                    return option.get("format")
+        return None
+
     # TODO: add sorts and filters
     prop_gettr = lambda name, d: (
         name,
         extract_types(d),
-        d.get("format"),
+        extract_format(d),
         hashable_default(d),
         d.get("deprecated"),
     )