1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements.  See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership.  The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License.  You may obtain a copy of the License at
8#
9#   http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied.  See the License for the
15# specific language governing permissions and limitations
16# under the License.
17import tvm
18
19def test_array():
20    a = tvm.convert([1,2,3])
21    assert len(a) == 3
22    assert a[-1].value == 3
23    a_slice = a[-3:-1]
24    assert (a_slice[0].value, a_slice[1].value) == (1, 2)
25
26def test_array_save_load_json():
27    a = tvm.convert([1,2,3])
28    json_str = tvm.save_json(a)
29    a_loaded = tvm.load_json(json_str)
30    assert(a_loaded[1].value == 2)
31
32
33def test_map():
34    a = tvm.var('a')
35    b = tvm.var('b')
36    amap = tvm.convert({a: 2,
37                        b: 3})
38    assert a in amap
39    assert len(amap) == 2
40    dd = dict(amap.items())
41    assert a in dd
42    assert b in dd
43    assert a + 1 not in amap
44
45
46def test_str_map():
47    amap = tvm.convert({'a': 2, 'b': 3})
48    assert 'a' in amap
49    assert len(amap) == 2
50    dd = dict(amap.items())
51    assert amap['a'].value == 2
52    assert 'a' in dd
53    assert 'b' in dd
54
55
56def test_map_save_load_json():
57    a = tvm.var('a')
58    b = tvm.var('b')
59    amap = tvm.convert({a: 2,
60                        b: 3})
61    json_str = tvm.save_json(amap)
62    amap = tvm.load_json(json_str)
63    assert len(amap) == 2
64    dd = {kv[0].name : kv[1].value for kv in amap.items()}
65    assert(dd == {"a": 2, "b": 3})
66
67
68def test_in_container():
69    arr = tvm.convert(['a', 'b', 'c'])
70    assert 'a' in arr
71    assert tvm.make.StringImm('a') in arr
72    assert 'd' not in arr
73
74if __name__ == "__main__":
75    test_str_map()
76    test_array()
77    test_map()
78    test_array_save_load_json()
79    test_map_save_load_json()
80    test_in_container()
81