## Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
##
## NVIDIA CORPORATION and its licensors retain all intellectual property
## and proprietary rights in and to this software, related documentation
## and any modifications thereto.  Any use, reproduction, disclosure or
## distribution of this software and related documentation without an express
## license agreement from NVIDIA CORPORATION is strictly prohibited.
##
import omni.kit.test
import omni.kit.app
from omni.kit.test.async_unittest import AsyncTestCase
import omni.usd

from omni.kit import ui_test
from pxr import Sdf, UsdShade
from omni.kit.test_suite.helpers import open_stage, get_test_data_path, select_prims, wait_stage_loading, wait_for_window, handle_assign_material_dialog, arrange_windows


class ViewportAssignMaterialSingle(AsyncTestCase):
    # Before running each test
    async def setUp(self):
        await arrange_windows()
        await open_stage(get_test_data_path(__name__, "bound_shapes.usda"))

    # After running each test
    async def tearDown(self):
        await wait_stage_loading()

    DATA = [["/World/Cube", "/World/Cone", "/World/Sphere", "/World/Cylinder"],
            ["", "/World/Looks/OmniGlass", "/World/Looks/OmniPBR", "/World/Looks/OmniSurface_Plastic"]]

    async def test_l1_viewport_assign_material_single(self):
        await ui_test.find("Stage").focus()
        await ui_test.find("Viewport").focus()

        usd_context = omni.usd.get_context()
        stage = usd_context.get_stage()

        for to_select in self.DATA[0]:
            for index, mtl_path in enumerate(self.DATA[1]):
                # select prim
                await select_prims([to_select])

                # wait for material to load & UI to refresh
                await wait_stage_loading()

                # right click on viewport
                await ui_test.find("Viewport").right_click()
                await ui_test.human_delay(50)

                # click context menu "Assign Material"
                await ui_test.select_context_menu("Assign Material")

                # use assign material dialog
                await handle_assign_material_dialog(index)

                # verify binding
                prim = stage.GetPrimAtPath(to_select)
                bound_material, _ = UsdShade.MaterialBindingAPI(prim).ComputeBoundMaterial()
                self.assertTrue(bound_material.GetPrim().GetPrimPath().pathString == mtl_path)
