# The MIT License (MIT)
# Copyright (c) Microsoft Corporation. All rights reserved.

import copy
import copy
import os
import unittest
import uuid
from unittest.mock import MagicMock

import pytest

import azure.cosmos.exceptions as exceptions
import test_config
from azure.cosmos import CosmosClient, PartitionKey
from azure.cosmos.database import DatabaseProxy

EGS_SKIP_REASON = "Set COSMOS_ENABLE_EGS_TESTS=1 to run; requires server-side EGS preview."



VectorPolicyTestData = {
    "valid_vector_indexing_policy" : {
        "indexing_policy": {
            "vectorIndexes": [
                {"path": "/vector1", "type": "flat"},
                {"path": "/vector2", "type": "quantizedFlat", "quantizerType": "product", "quantizationByteSize": 8},
                {"path": "/vector3", "type": "diskANN", "quantizerType": "product", "quantizationByteSize": 8,
                 "vectorIndexShardKey": ["/city"], "indexingSearchListSize": 50},
                {"path": "/vector4", "type": "diskANN", "quantizerType": "spherical", "indexingSearchListSize": 50},
            ]
        },
        "vector_embedding_policy": {
            "vectorEmbeddings": [
                {
                    "path": "/vector1",
                    "dataType": "float32",
                    "dimensions": 256,
                    "distanceFunction": "euclidean"
                },
                {
                    "path": "/vector2",
                    "dataType": "int8",
                    "dimensions": 200,
                    "distanceFunction": "dotproduct"
                },
                {
                    "path": "/vector3",
                    "dataType": "uint8",
                    "dimensions": 400,
                    "distanceFunction": "cosine"
                },
                {
                    "path": "/vector4",
                    "dataType": "uint8",
                    "dimensions": 400,
                    "distanceFunction": "euclidean"
                },
            ]
        }
    },
    "valid_vector_embedding_policy_with_source": {
        "vectorEmbeddings": [
            {
                "path": "/embedding",
                "dataType": "float32",
                "dimensions": 1536,
                "distanceFunction": "cosine",
                "embeddingSource": {
                    "sourcePaths": [
                        "/journal_title",
                        "/title",
                        "/toc_abstract",
                        "/abstract",
                        "/full_text",
                    ],
                    "deploymentName": "text-embedding-3-small",
                    "modelName": "text-embedding-3-small",
                    "endpoint": "https://example.com",
                    "authType": "ApiKey",
                },
            }
        ]
    },
}

@pytest.mark.cosmosSearchQuery
class TestVectorPolicy(unittest.TestCase):
    client: CosmosClient = None
    key_client: CosmosClient = None
    data_client: CosmosClient = None
    host = test_config.TestConfig.host
    masterKey = test_config.TestConfig.masterKey
    connectionPolicy = test_config.TestConfig.connectionPolicy

    @classmethod
    def setUpClass(cls):
        if (cls.masterKey == '[YOUR_KEY_HERE]' or
                cls.host == '[YOUR_ENDPOINT_HERE]'):
            raise Exception(
                "You must specify your Azure Cosmos account values for "
                "'masterKey' and 'host' at the top of this class to run the "
                "tests.")

        cls.client = CosmosClient(cls.host, cls.masterKey)
        cls.key_client = cls.client  # alias  -  control-plane operations stay on key-auth (Batch 17 prep)
        # AAD data client added for parity with the key/data client setup. Not exercised
        # here because every runnable test in this file is control-plane (vector indexing/
        # embedding policy validation via create_container / replace_container / read).
        # When per-test data-plane operations are added (e.g., vector similarity queries
        # against a populated container), route those through cls.data_client.
        cls.data_client = test_config.TestConfig.create_data_client()
        cls.created_database = cls.client.get_database_client(test_config.TestConfig.TEST_DATABASE_ID)
        cls.test_db = cls.client.create_database(str(uuid.uuid4()))

    @classmethod
    def tearDownClass(cls):
        test_config.TestConfig.try_delete_database_with_id(cls.client, cls.test_db.id)

    def test_create_valid_vector_embedding_policy(self):
        # Using valid data types
        data_types = ["float32", "float16", "int8", "uint8"]
        for data_type in data_types:
            vector_embedding_policy = {
                "vectorEmbeddings": [
                    {
                        "path": "/vector1",
                        "dataType": data_type,
                        "dimensions": 256,
                        "distanceFunction": "euclidean"
                    }]}
            container_id = 'vector_container_' + data_type
            created_container = self.test_db.create_container(
                id=container_id,
                partition_key=PartitionKey(path="/id"),
                vector_embedding_policy=vector_embedding_policy)
            try:
                properties = created_container.read()
                assert properties["vectorEmbeddingPolicy"]["vectorEmbeddings"][0]["dataType"] == data_type
            finally:
                self.test_db.delete_container(container_id)

    @unittest.skip
    def test_create_valid_vector_indexing_policy(self):
        test_data = VectorPolicyTestData["valid_vector_indexing_policy"]
        indexing_policy = test_data["indexing_policy"]
        vector_embedding_policy = test_data["vector_embedding_policy"]

        created_container = self.test_db.create_container(
            id="container_" + str(uuid.uuid4()),
            partition_key=PartitionKey(path="/id"),
            vector_embedding_policy=vector_embedding_policy,
            indexing_policy=indexing_policy)
        try:
            properties = created_container.read()
            assert properties['indexingPolicy']['vectorIndexes'] == indexing_policy['vectorIndexes']
        finally:
            self.test_db.delete_container(created_container.id)

    def test_create_vector_embedding_container(self):
        indexing_policy = {
            "vectorIndexes": [
                {"path": "/vector1", "type": "flat"},
                {"path": "/vector2", "type": "quantizedFlat", "quantizerType": "product", "quantizationByteSize": 8},
                {"path": "/vector3", "type": "diskANN", "quantizerType": "product", "quantizationByteSize": 8, "vectorIndexShardKey": ["/city"], "indexingSearchListSize": 50}
            ]
        }
        vector_embedding_policy = {
            "vectorEmbeddings": [
                {
                    "path": "/vector1",
                    "dataType": "float32",
                    "dimensions": 256,
                    "distanceFunction": "euclidean"
                },
                {
                    "path": "/vector2",
                    "dataType": "int8",
                    "dimensions": 200,
                    "distanceFunction": "dotproduct"
                },
                {
                    "path": "/vector3",
                    "dataType": "uint8",
                    "dimensions": 400,
                    "distanceFunction": "cosine"
                }
            ]
        }
        container_id = "vector_container" + str(uuid.uuid4())
        created_container = self.test_db.create_container(
            id=container_id,
            partition_key=PartitionKey(path="/id"),
            vector_embedding_policy=vector_embedding_policy,
            indexing_policy=indexing_policy
        )
        try:
            properties = created_container.read()
            assert properties["vectorEmbeddingPolicy"] == vector_embedding_policy
            assert properties["indexingPolicy"]["vectorIndexes"] == indexing_policy["vectorIndexes"]
        finally:
            self.test_db.delete_container(container_id)

        # Pass a vector indexing policy with hierarchical vectorIndexShardKey value
        indexing_policy = {
            "vectorIndexes": [
                {"path": "/vector2", "type": "diskANN", "quantizerType": "product", 'quantizationByteSize': 64, 'indexingSearchListSize': 100, "vectorIndexShardKey": ["/country/city"]}]
        }
        container_id = "vector_container" + str(uuid.uuid4())
        created_container = self.test_db.create_container(
            id=container_id,
            partition_key=PartitionKey(path="/id"),
            indexing_policy=indexing_policy,
            vector_embedding_policy=vector_embedding_policy
        )
        try:
            properties = created_container.read()
            assert properties["vectorEmbeddingPolicy"] == vector_embedding_policy
            assert properties["indexingPolicy"]["vectorIndexes"] == indexing_policy["vectorIndexes"]
        finally:
            self.test_db.delete_container(container_id)

    def test_replace_vector_indexing_policy(self):
        # Replace should work so long as the new indexing policy doesn't change the vector indexes, and as long as
        # the previously defined vector embedding policy is also provided.
        vector_embedding_policy = {
            "vectorEmbeddings": [
                {
                    "path": "/vector1",
                    "dataType": "float32",
                    "dimensions": 256,
                    "distanceFunction": "euclidean"
                }
            ]
        }
        indexing_policy = {
            "indexingMode": "consistent",
            "automatic": True,
            "includedPaths": [
                {
                    "path": "/*"
                }
            ],
            "excludedPaths": [
                {
                    "path": "/vector1/*"
                },
                {
                    "path": "/\"_etag\"/?"
                }
            ],
            "fullTextIndexes": [],
            "vectorIndexes": [
                {
                    "path": "/vector1",
                    "type": "diskANN",
                    "quantizerType": "product",
                    "quantizationByteSize": 128,
                    "indexingSearchListSize": 100
                }
            ]
        }
        container_id = "vector_container" + str(uuid.uuid4())
        created_container = self.test_db.create_container(
            id=container_id,
            partition_key=PartitionKey(path="/id"),
            indexing_policy=indexing_policy,
            vector_embedding_policy=vector_embedding_policy
        )
        new_indexing_policy = {
            "indexingMode": "consistent",
            "automatic": True,
            "includedPaths": [
                {"path": "/color/?"},
                {"path": "/description/?"},
                {"path": "/cost/?"}
            ],
            "excludedPaths": [
                {"path": "/*"},
                {"path": "/vector1/*"},
                {"path": "/\"_etag\"/?"}
            ],
            "fullTextIndexes": [],
            "vectorIndexes": [
                {
                    "path": "/vector1",
                    "type": "diskANN",
                    "quantizerType": "product",
                    "quantizationByteSize": 128,
                    "indexingSearchListSize": 100
                }]
        }

        try:
            self.test_db.replace_container(
                created_container,
                PartitionKey(path="/id"),
                vector_embedding_policy=vector_embedding_policy,
                indexing_policy=new_indexing_policy)
            properties = created_container.read()
            assert properties["vectorEmbeddingPolicy"] == vector_embedding_policy
            assert properties["indexingPolicy"]["vectorIndexes"] == indexing_policy["vectorIndexes"]
        finally:
            self.test_db.delete_container(container_id)

    def test_fail_create_vector_indexing_policy(self):
        vector_embedding_policy = {
            "vectorEmbeddings": [
                {
                    "path": "/vector1",
                    "dataType": "float32",
                    "dimensions": 256,
                    "distanceFunction": "euclidean"
                },
                {
                    "path": "/vector2",
                    "dataType": "int8",
                    "dimensions": 200,
                    "distanceFunction": "dotproduct"
                }
            ]
        }

        # Pass a vector indexing policy without embedding policy
        indexing_policy = {
            "vectorIndexes": [
                {"path": "/vector1", "type": "flat"}]
        }
        try:
            self.test_db.create_container(
                id='vector_container',
                partition_key=PartitionKey(path="/id"),
                indexing_policy=indexing_policy
            )
            pytest.fail("Container creation should have failed for lack of embedding policy.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "vector1 not matching in Embedding's path" in e.http_error_message

        # Pass a vector indexing policy with an invalid type
        indexing_policy = {
            "vectorIndexes": [
                {"path": "/vector1", "type": "notFlat"}]
        }
        try:
            self.test_db.create_container(
                id='vector_container',
                partition_key=PartitionKey(path="/id"),
                indexing_policy=indexing_policy,
                vector_embedding_policy=vector_embedding_policy
            )
            pytest.fail("Container creation should have failed for wrong index type.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "Index Type::notFlat is invalid" in e.http_error_message

        # Pass a vector indexing policy with non-matching path
        indexing_policy = {
            "vectorIndexes": [
                {"path": "/vector3", "type": "flat"}]
        }
        try:
            self.test_db.create_container(
                id='vector_container',
                partition_key=PartitionKey(path="/id"),
                indexing_policy=indexing_policy,
                vector_embedding_policy=vector_embedding_policy
            )
            pytest.fail("Container creation should have failed for index mismatch.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "vector3 not matching in Embedding's path" in e.http_error_message

        # Pass a vector indexing policy with wrong quantizationByteSize value
        indexing_policy = {
            "vectorIndexes": [
                {"path": "/vector1", "type": "quantizedFlat", "quantizationByteSize": 0}]
        }
        try:
            self.test_db.create_container(
                id='vector_container',
                partition_key=PartitionKey(path="/id"),
                indexing_policy=indexing_policy,
                vector_embedding_policy=vector_embedding_policy
            )
            pytest.fail("Container creation should have failed for value mismatch.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "The Vector Indexing Policy parameter QuantizationByteSize value :: 0 is out of range. The allowed range is between 1 and 256."\
                   in e.http_error_message

        # Pass a vector indexing policy with wrong indexingSearchListSize value
        indexing_policy = {
            "vectorIndexes": [
                {"path": "/vector1", "type": "diskANN", "indexingSearchListSize": 5}]
        }
        try:
            self.test_db.create_container(
                id='vector_container',
                partition_key=PartitionKey(path="/id"),
                indexing_policy=indexing_policy,
                vector_embedding_policy=vector_embedding_policy
            )
            pytest.fail("Container creation should have failed for value mismatch.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "IndexingSearchListSize value :: 5 is out of range. The allowed range is between 25 and 500."\
                   in e.http_error_message

        # Pass a vector indexing policy with wrong vectorIndexShardKey value
        indexing_policy = {
            "vectorIndexes": [
                {"path": "/vector2", "type": "diskANN", "vectorIndexShardKey": ["country"]}]
        }
        try:
            self.test_db.create_container(
                id='vector_container',
                partition_key=PartitionKey(path="/id"),
                indexing_policy=indexing_policy,
                vector_embedding_policy=vector_embedding_policy
            )
            pytest.fail("Container creation should have failed for value mismatch.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "The Vector Indexing Policy has an invalid Shard Path: country." in e.http_error_message

        # Pass a vector indexing policy with too many shard paths
        indexing_policy = {
            "vectorIndexes": [
                {"path": "/vector2", "type": "diskANN", "vectorIndexShardKey": ["/country", "/city", "/zipcode"]}]
        }
        try:
            self.test_db.create_container(
                id='vector_container',
                partition_key=PartitionKey(path="/id"),
                indexing_policy=indexing_policy,
                vector_embedding_policy=vector_embedding_policy
            )
            pytest.fail("Container creation should have failed for value mismatch.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "The number of shard paths defined in the Vector Indexing Policy: 3 exceeds the maximum: 1." \
                   in e.http_error_message

        # Pass a vector indexing policy with an invalid type for vectorIndexShardKey
        indexing_policy = {
            "vectorIndexes": [
                {"path": "/vector2", "type": "diskANN", "vectorIndexShardKey": "/country"}]
        }
        try:
            self.test_db.create_container(
                id='vector_container',
                partition_key=PartitionKey(path="/id"),
                indexing_policy=indexing_policy,
                vector_embedding_policy=vector_embedding_policy
            )
            pytest.fail("Container creation should have failed for value mismatch.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "One of the specified inputs is invalid" \
                   in e.http_error_message

        # Pass a vector indexing policy with dimensions above 512 to test the max range of  "quantizationByteSize"
        vector_embedding_policy = {
            "vectorEmbeddings": [
                {
                    "path": "/vector1",
                    "dataType": "float32",
                    "dimensions": 550,
                    "distanceFunction": "euclidean"
                }]}
        indexing_policy = {
            "vectorIndexes": [
                {"path": "/vector1", "type": "quantizedFlat", "quantizationByteSize": 513 }]
        }
        try:
            self.test_db.create_container(
                id='vector_container',
                partition_key=PartitionKey(path="/id"),
                indexing_policy=indexing_policy,
                vector_embedding_policy=vector_embedding_policy
            )
            pytest.fail("Container creation should have failed for value mismatch.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "The Vector Indexing Policy parameter QuantizationByteSize value :: 513 is out of range." \
                   in e.http_error_message

    def test_fail_replace_vector_indexing_policy(self):
        vector_embedding_policy = {
            "vectorEmbeddings": [
                {
                    "path": "/vector1",
                    "dataType": "float32",
                    "dimensions": 256,
                    "distanceFunction": "euclidean"
                }]}
        indexing_policy = {
            "vectorIndexes": [
                {"path": "/vector1", "type": "flat"}]
        }
        container_id = "vector_container" + str(uuid.uuid4())
        created_container = self.test_db.create_container(
            id=container_id,
            partition_key=PartitionKey(path="/id"),
            indexing_policy=indexing_policy,
            vector_embedding_policy=vector_embedding_policy
        )
        try:
            # don't provide vector embedding policy
            try:
                self.test_db.replace_container(
                    created_container,
                    PartitionKey(path="/id"),
                    indexing_policy=indexing_policy)
                pytest.fail("Container replace should have failed for missing embedding policy.")
            except exceptions.CosmosHttpResponseError as e:
                assert e.status_code == 400
                assert ("The Vector Indexing Policy's path::/vector1 not matching in Embedding's path."
                        in e.http_error_message)
            # using a new indexing policy
            new_indexing_policy = {
                "vectorIndexes": [
                    {"path": "/vector1", "type": "quantizedFlat"}]
            }
            try:
                self.test_db.replace_container(
                    created_container,
                    PartitionKey(path="/id"),
                    vector_embedding_policy=vector_embedding_policy,
                    indexing_policy=new_indexing_policy)
                pytest.fail("Container replace should have failed for new indexing policy.")
            except exceptions.CosmosHttpResponseError as e:
                assert e.status_code == 400
                assert ("Paths in existing vector indexing policy cannot be modified in Collection Replace"
                        in e.http_error_message)
            # using a new vector embedding policy
            new_embedding_policy = {
                "vectorEmbeddings": [
                    {
                        "path": "/vector1",
                        "dataType": "float32",
                        "dimensions": 384,
                        "distanceFunction": "euclidean"}]}
            try:
                self.test_db.replace_container(
                    created_container,
                    PartitionKey(path="/id"),
                    vector_embedding_policy=new_embedding_policy,
                    indexing_policy=indexing_policy)
                pytest.fail("Container replace should have failed for new embedding policy.")
            except exceptions.CosmosHttpResponseError as e:
                assert e.status_code == 400
                assert ("Paths in existing embedding policy cannot be modified in Collection Replace"
                        in e.http_error_message)
        finally:
            self.test_db.delete_container(container_id)

    def test_fail_create_vector_embedding_policy(self):
        # Using invalid data type
        vector_embedding_policy = {
            "vectorEmbeddings": [
                {
                    "path": "/vector1",
                    "dataType": "float33",
                    "dimensions": 256,
                    "distanceFunction": "euclidean"
                }]}
        try:
            self.test_db.create_container(
                id='vector_container',
                partition_key=PartitionKey(path="/id"),
                vector_embedding_policy=vector_embedding_policy)
            pytest.fail("Container creation should have failed but succeeded.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "Vector Embedding Policy has an invalid DataType" in e.http_error_message

        # Using too many dimensions
        vector_embedding_policy = {
            "vectorEmbeddings": [
                {
                    "path": "/vector1",
                    "dataType": "float32",
                    "dimensions": 8000,
                    "distanceFunction": "euclidean"
                }]}
        try:
            self.test_db.create_container(
                id='vector_container',
                partition_key=PartitionKey(path="/id"),
                vector_embedding_policy=vector_embedding_policy)
            pytest.fail("Container creation should have failed but succeeded.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "Vector Embedding Policy has Dimensions:8000 which is more than the maximum" \
                   " supported value" in e.http_error_message

        # Using negative dimensions
        vector_embedding_policy = {
            "vectorEmbeddings": [
                {
                    "path": "/vector1",
                    "dataType": "float32",
                    "dimensions": -1,
                    "distanceFunction": "euclidean"
                }]}
        try:
            self.test_db.create_container(
                id='vector_container',
                partition_key=PartitionKey(path="/id"),
                vector_embedding_policy=vector_embedding_policy)
            pytest.fail("Container creation should have failed but succeeded.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "The Vector Embedding Policy has Dimensions:-1" in e.http_error_message

        # Using invalid distance function
        vector_embedding_policy = {
            "vectorEmbeddings": [
                {
                    "path": "/vector1",
                    "dataType": "float32",
                    "dimensions": 256,
                    "distanceFunction": "handMeasured"
                }]}
        try:
            self.test_db.create_container(
                id='vector_container',
                partition_key=PartitionKey(path="/id"),
                vector_embedding_policy=vector_embedding_policy
            )
            pytest.fail("Container creation should have failed but succeeded.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "The Vector Embedding Policy has an invalid DistanceFunction:handMeasured" in e.http_error_message

    @unittest.skipUnless(os.getenv("COSMOS_ENABLE_EGS_TESTS"), EGS_SKIP_REASON)
    def test_create_vector_embedding_with_embedding_source(self):
        for auth_type in ["ApiKey", "Entra"]:
            vector_embedding_policy = copy.deepcopy(
                VectorPolicyTestData["valid_vector_embedding_policy_with_source"])
            vector_embedding_policy["vectorEmbeddings"][0]["embeddingSource"]["authType"] = auth_type
            container_id = "vector_embedding_source_container_" + auth_type.lower() + "_" + str(uuid.uuid4())
            created_container = self.test_db.create_container(
                id=container_id,
                partition_key=PartitionKey(path="/id"),
                vector_embedding_policy=vector_embedding_policy,
            )
            try:
                properties = created_container.read()
                assert properties["vectorEmbeddingPolicy"] == vector_embedding_policy
            finally:
                self.test_db.delete_container(container_id)

    @unittest.skipUnless(os.getenv("COSMOS_ENABLE_EGS_TESTS"), EGS_SKIP_REASON)
    def test_fail_create_vector_embedding_source_invalid_auth_type(self):
        vector_embedding_policy = copy.deepcopy(
            VectorPolicyTestData["valid_vector_embedding_policy_with_source"])
        vector_embedding_policy["vectorEmbeddings"][0]["embeddingSource"]["authType"] = "NotAnAuthType"

        try:
            self.test_db.create_container(
                id="vector_embedding_source_bad_auth_" + str(uuid.uuid4()),
                partition_key=PartitionKey(path="/id"),
                vector_embedding_policy=vector_embedding_policy,
            )
            pytest.fail("Container creation should have failed for invalid embeddingSource authType.")
        except exceptions.CosmosHttpResponseError as e:
            assert e.status_code == 400
            assert "The auth type provided in the embedding source of vector policy is invalid." \
                in e.http_error_message

    @unittest.skipUnless(os.getenv("COSMOS_ENABLE_EGS_TESTS"), EGS_SKIP_REASON)
    def test_fail_create_vector_embedding_source_missing_required_fields(self):
        for field_name in ["sourcePaths", "endpoint", "deploymentName"]:
            vector_embedding_policy = copy.deepcopy(
                VectorPolicyTestData["valid_vector_embedding_policy_with_source"])
            embedding_source = vector_embedding_policy["vectorEmbeddings"][0]["embeddingSource"]
            if field_name == "sourcePaths":
                embedding_source["sourcePaths"] = []
            else:
                del embedding_source[field_name]

            container_id = "vector_embedding_source_missing_" + field_name.lower() + "_" + str(uuid.uuid4())
            try:
                self.test_db.create_container(
                    id=container_id,
                    partition_key=PartitionKey(path="/id"),
                    vector_embedding_policy=vector_embedding_policy,
                )
                pytest.fail(
                    "Container creation should have failed for embeddingSource missing '{}'.".format(field_name)
                )
            except exceptions.CosmosHttpResponseError as e:
                assert e.status_code == 400, (
                    "Expected 400 for missing '{}', got {}".format(field_name, e.status_code)
                )


if __name__ == '__main__':
    unittest.main()


class TestVectorEmbeddingPolicyOffline(unittest.TestCase):

    def test_embedding_source_round_trips_through_create_container_definition(self):
        vector_embedding_policy = VectorPolicyTestData["valid_vector_embedding_policy_with_source"]

        mock_connection = MagicMock()
        mock_connection.CreateContainer.return_value = {"id": "c", "_self": "dbs/x/colls/c"}
        db = DatabaseProxy(mock_connection, "x")

        db.create_container(
            id="c",
            partition_key=PartitionKey(path="/id"),
            vector_embedding_policy=vector_embedding_policy,
        )

        sent_collection = mock_connection.CreateContainer.call_args.kwargs["collection"]
        assert sent_collection["vectorEmbeddingPolicy"] == vector_embedding_policy
