# Copyright 2018 The dm_control Authors.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ============================================================================

"""Tests for `dm_control.mjcf.copier`."""

import os

from absl.testing import absltest

from dm_control import mjcf

from dm_control.mjcf import parser

import numpy as np

_ASSETS_DIR = os.path.join(os.path.dirname(__file__), 'test_assets')

_TEST_MODEL_XML = os.path.join(_ASSETS_DIR, 'test_model.xml')

_MODEL_WITH_ASSETS_XML = os.path.join(_ASSETS_DIR, 'model_with_assets.xml')

class CopierTest(absltest.TestCase):

def testSimpleCopy(self):

mjcf_model = parser.from_path(_TEST_MODEL_XML)

mixin = mjcf.RootElement(model='test_mixin')

mixin.compiler.boundmass = 1

mjcf_model.include_copy(mixin)

self.assertEqual(mjcf_model.model, 'test') # Model name should not change

self.assertEqual(mjcf_model.compiler.boundmass, mixin.compiler.boundmass)

mixin.compiler.boundinertia = 2

mjcf_model.include_copy(mixin)

self.assertEqual(mjcf_model.compiler.boundinertia,

mixin.compiler.boundinertia)

mixin.compiler.boundinertia = 1

with self.assertRaisesRegex(ValueError, 'Conflicting values'):

mjcf_model.include_copy(mixin)

mixin.worldbody.add('body', name='b_0', pos=[0, 1, 2])

mjcf_model.include_copy(mixin, override_attributes=True)

self.assertEqual(mjcf_model.compiler.boundmass, mixin.compiler.boundmass)

self.assertEqual(mjcf_model.compiler.boundinertia,

mixin.compiler.boundinertia)

np.testing.assert_array_equal(mjcf_model.worldbody.body['b_0'].pos,

[0, 1, 2])

def testCopyingWithReference(self):

sensor_mixin = mjcf.RootElement('sensor_mixin')

touch_site = sensor_mixin.worldbody.add('site', name='touch_site')

sensor_mixin.sensor.add('touch', name='touch_sensor', site=touch_site)

mjcf_model = mjcf.RootElement('model')

mjcf_model.include_copy(sensor_mixin)

# Copied reference should be updated to the copied site.

self.assertIs(mjcf_model.find('sensor', 'touch_sensor').site,

mjcf_model.find('site', 'touch_site'))

def testCopyingWithAssets(self):

mjcf_model = parser.from_path(_MODEL_WITH_ASSETS_XML)

copied = mjcf.RootElement()

copied.include_copy(mjcf_model)

original_assets = (mjcf_model.find_all('mesh')

+ mjcf_model.find_all('texture')

+ mjcf_model.find_all('hfield'))

copied_assets = (copied.find_all('mesh')

+ copied.find_all('texture')

+ copied.find_all('hfield'))

self.assertLen(copied_assets, len(original_assets))

for original_asset, copied_asset in zip(original_assets, copied_assets):

self.assertIs(copied_asset.file, original_asset.file)

if __name__ == '__main__':

absltest.main()