wehub-resource-sync
121 lines
4.0 KiB
Python
121 lines
4.0 KiB
Python
from mlflow.models.auth_policy import AuthPolicy, SystemAuthPolicy, UserAuthPolicy
|
|
from mlflow.models.resources import (
|
|
DatabricksFunction,
|
|
DatabricksServingEndpoint,
|
|
DatabricksUCConnection,
|
|
DatabricksVectorSearchIndex,
|
|
)
|
|
|
|
|
|
def test_complete_auth_policy():
|
|
system_auth_policy = SystemAuthPolicy(
|
|
resources=[
|
|
DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
|
|
DatabricksVectorSearchIndex(index_name="rag.studio_bugbash.databricks_docs_index"),
|
|
DatabricksFunction(function_name="rag.studio.test_function_a"),
|
|
DatabricksUCConnection(connection_name="test_connection_1"),
|
|
]
|
|
)
|
|
|
|
user_auth_policy = UserAuthPolicy(
|
|
api_scopes=[
|
|
"catalog.catalogs",
|
|
"vectorsearch.vector-search-indexes",
|
|
"workspace.workspace",
|
|
]
|
|
)
|
|
|
|
auth_policy = AuthPolicy(
|
|
user_auth_policy=user_auth_policy, system_auth_policy=system_auth_policy
|
|
)
|
|
|
|
serialized_auth_policy = auth_policy.to_dict()
|
|
|
|
expected_serialized_auth_policy = {
|
|
"user_auth_policy": {
|
|
"api_scopes": [
|
|
"catalog.catalogs",
|
|
"vectorsearch.vector-search-indexes",
|
|
"workspace.workspace",
|
|
]
|
|
},
|
|
"system_auth_policy": {
|
|
"resources": {
|
|
"databricks": {
|
|
"serving_endpoint": [{"name": "databricks-mixtral-8x7b-instruct"}],
|
|
"vector_search_index": [{"name": "rag.studio_bugbash.databricks_docs_index"}],
|
|
"function": [{"name": "rag.studio.test_function_a"}],
|
|
"uc_connection": [{"name": "test_connection_1"}],
|
|
},
|
|
"api_version": "1",
|
|
}
|
|
},
|
|
}
|
|
assert serialized_auth_policy == expected_serialized_auth_policy
|
|
|
|
|
|
def test_user_auth_policy():
|
|
user_auth_policy = UserAuthPolicy(
|
|
api_scopes=[
|
|
"catalog.catalogs",
|
|
"vectorsearch.vector-search-indexes",
|
|
"workspace.workspace",
|
|
]
|
|
)
|
|
|
|
auth_policy = AuthPolicy(user_auth_policy=user_auth_policy)
|
|
|
|
serialized_auth_policy = auth_policy.to_dict()
|
|
|
|
expected_serialized_auth_policy = {
|
|
"system_auth_policy": {},
|
|
"user_auth_policy": {
|
|
"api_scopes": [
|
|
"catalog.catalogs",
|
|
"vectorsearch.vector-search-indexes",
|
|
"workspace.workspace",
|
|
]
|
|
},
|
|
}
|
|
assert serialized_auth_policy == expected_serialized_auth_policy
|
|
|
|
|
|
def test_system_auth_policy():
|
|
system_auth_policy = SystemAuthPolicy(
|
|
resources=[
|
|
DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
|
|
DatabricksVectorSearchIndex(index_name="rag.studio_bugbash.databricks_docs_index"),
|
|
DatabricksFunction(function_name="rag.studio.test_function_a"),
|
|
DatabricksUCConnection(connection_name="test_connection_1"),
|
|
]
|
|
)
|
|
|
|
auth_policy = AuthPolicy(system_auth_policy=system_auth_policy)
|
|
|
|
serialized_auth_policy = auth_policy.to_dict()
|
|
|
|
expected_serialized_auth_policy = {
|
|
"system_auth_policy": {
|
|
"resources": {
|
|
"databricks": {
|
|
"serving_endpoint": [{"name": "databricks-mixtral-8x7b-instruct"}],
|
|
"vector_search_index": [{"name": "rag.studio_bugbash.databricks_docs_index"}],
|
|
"function": [{"name": "rag.studio.test_function_a"}],
|
|
"uc_connection": [{"name": "test_connection_1"}],
|
|
},
|
|
"api_version": "1",
|
|
}
|
|
},
|
|
"user_auth_policy": {},
|
|
}
|
|
assert serialized_auth_policy == expected_serialized_auth_policy
|
|
|
|
|
|
def test_empty_auth_policy():
|
|
auth_policy = AuthPolicy()
|
|
|
|
serialized_auth_policy = auth_policy.to_dict()
|
|
|
|
expected_serialized_auth_policy = {"system_auth_policy": {}, "user_auth_policy": {}}
|
|
assert serialized_auth_policy == expected_serialized_auth_policy
|