[WORKFLOW] AI update based on proto changes from commit c771880

Author: viambot

src/viam/app/app_client.py

@@ -756,6 +756,7 @@ async def update_organization(
             public_namespace (Optional[str]): If provided, sets the org's namespace if it hasn't already been set.
             region (Optional[str]): If provided, updates the org's region.
             cid (Optional[str]): If provided, update's the org's CRM ID.
+            default_fragments (Optional[FragmentImportList]): If provided, updates the org's default fragments.
 
         Raises:
             GRPCError: If the org's namespace has already been set, or if the provided namespace is already taken.

tests/mocks/services.py

@@ -7,7 +7,8 @@
 from numpy.typing import NDArray
 
 from viam.app.data_client import DataClient
-from viam.gen.app.v1.app_pb2 import FragmentHistoryEntry, GetFragmentHistoryRequest, GetFragmentHistoryResponse
+from viam.gen.app.data import FragmentHistoryEntry, GetFragmentHistoryRequest, GetFragmentHistoryResponse
+from viam.gen.app.packages import FragmentImportList, PackageType
 from viam.media.video import ViamImage
 from viam.proto.app import (
     AddRoleRequest,
@@ -314,8 +315,6 @@
     TrainingJobMetadata,
     UnimplementedMLTrainingServiceBase,
 )
-from viam.proto.app.packages import PackageType
-from viam.proto.app.robot import ComponentConfig
 from viam.proto.common import (
     DoCommandRequest,
     DoCommandResponse,
@@ -1388,6 +1387,7 @@ def __init__(
         api_keys_with_authorizations: List[APIKeyWithAuthorizations],
         items: List[RegistryItem],
         package_type: PackageType.ValueType,
+        default_fragments_for_org: Optional[FragmentImportList] = None,
     ):
         self.organizations = organizations
         self.location = location
@@ -1411,6 +1411,7 @@ def __init__(
         self.api_keys_with_authorizations = api_keys_with_authorizations
         self.items = items
         self.package_type = package_type
+        self.default_fragments_for_org = default_fragments_for_org
         self.send_email_invite = False
         self.organization_metadata = {}
         self.location_metadata = {}
@@ -1440,7 +1441,7 @@ async def ListOrganizationsByUser(self, stream: Stream[ListOrganizationsByUserRe
     async def GetOrganization(self, stream: Stream[GetOrganizationRequest, GetOrganizationResponse]) -> None:
         request = await stream.recv_message()
         assert request is not None
-        await stream.send_message(GetOrganizationResponse(organization=self.organizations[0]))
+        await stream.send_message(GetOrganizationResponse(organization=Organization(id=self.id, name=self.name, created_on=self.organizations[0].created_on, public_namespace=self.organizations[0].public_namespace, default_region=self.organizations[0].default_region, cid=self.organizations[0].cid, default_fragments=self.default_fragments_for_org)))
 
     async def GetOrganizationNamespaceAvailability(
         self, stream: Stream[GetOrganizationNamespaceAvailabilityRequest, GetOrganizationNamespaceAvailabilityResponse]
@@ -1457,7 +1458,8 @@ async def UpdateOrganization(self, stream: Stream[UpdateOrganizationRequest, Upd
         self.update_cid = request.cid
         self.update_name = request.name
         self.update_namespace = request.public_namespace
-        await stream.send_message(UpdateOrganizationResponse(organization=self.organizations[0]))
+        self.update_default_fragments = request.default_fragments
+        await stream.send_message(UpdateOrganizationResponse(organization=Organization(id=self.id, name=self.name, created_on=self.organizations[0].created_on, public_namespace=self.update_namespace, default_region=self.update_region, cid=self.update_cid, default_fragments=self.update_default_fragments)))
 
     async def DeleteOrganization(self, stream: Stream[DeleteOrganizationRequest, DeleteOrganizationResponse]) -> None:
         request = await stream.recv_message()

tests/test_app_client.py

@@ -208,6 +208,8 @@
 MODULE_FILE_INFO = ModuleFileInfo(module_id=ID, version=VERSION, platform=PLATFORM)
 FILE = b"file"
 USER_DEFINED_METADATA = {"number": 0, "string": "string"}
+FRAGMENT_IMPORT = FragmentImport(fragment_id=ID, version=VERSION)
+FRAGMENT_IMPORT_LIST = FragmentImportList(fragments=[FRAGMENT_IMPORT])
 
 
 @pytest.fixture(scope="function")
@@ -270,6 +272,7 @@ async def test_get_organization(self, service: MockApp):
             client = AppClient(channel, METADATA)
             org = await client.get_organization(org_id=ID)
             assert org == ORGANIZATION
+            assert org.default_fragments == FRAGMENT_IMPORT_LIST
             available = await client.get_organization_namespace_availability(public_namespace=NAMESPACE)
             assert available == AVAILABLE
             assert service.namespace == NAMESPACE
@@ -297,12 +300,13 @@ async def test_list_organizations(self, service: MockApp):
     async def test_update_organization(self, service: MockApp):
         async with ChannelFor([service]) as channel:
             client = AppClient(channel, METADATA)
-            org = await client.update_organization(org_id=ID, name=NAME, public_namespace=PUBLIC_NAMESPACE, region=DEFAULT_REGION, cid=CID)
+            org = await client.update_organization(org_id=ID, name=NAME, public_namespace=PUBLIC_NAMESPACE, region=DEFAULT_REGION, cid=CID, default_fragments=FRAGMENT_IMPORT_LIST)
             assert org == ORGANIZATION
             assert service.update_region == DEFAULT_REGION
             assert service.update_cid == CID
             assert service.update_name == NAME
             assert service.update_namespace == PUBLIC_NAMESPACE
+            assert service.update_default_fragments == FRAGMENT_IMPORT_LIST
 
     async def test_update_organization_invite_authorizations(self, service: MockApp):
         async with ChannelFor([service]) as channel: