1+ """
2+ Tests for the parse_teams module.
3+ """
4+
5+ # Copyright (C) 2023 The Chronon Authors.
6+ #
7+ # Licensed under the Apache License, Version 2.0 (the "License");
8+ # you may not use this file except in compliance with the License.
9+ # You may obtain a copy of the License at
10+ #
11+ # http://www.apache.org/licenses/LICENSE-2.0
12+ #
13+ # Unless required by applicable law or agreed to in writing, software
14+ # distributed under the License is distributed on an "AS IS" BASIS,
15+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+ # See the License for the specific language governing permissions and
17+ # limitations under the License.
18+
19+ from ai .chronon .api .ttypes import GroupBy , Join , JoinPart , LabelParts , MetaData , Team
20+ from ai .chronon .cli .compile import parse_teams
21+
22+
23+ def test_update_metadata_with_existing_output_namespace ():
24+ """Test that update_metadata doesn't override existing outputNamespace."""
25+ # Setup
26+ team_name = "test_team"
27+ team_dict = {
28+ "default" : Team (outputNamespace = "default_namespace" ),
29+ team_name : Team (outputNamespace = "team_namespace" ),
30+ }
31+
32+ # Test with existing outputNamespace
33+ existing_namespace = "existing_namespace"
34+ obj = GroupBy (metaData = MetaData (
35+ team = team_name ,
36+ name = "test.group_by.name" ,
37+ outputNamespace = existing_namespace
38+ ))
39+
40+ # Call the function
41+ parse_teams .update_metadata (obj , team_dict )
42+
43+ # Verify outputNamespace wasn't changed
44+ assert obj .metaData .outputNamespace == existing_namespace
45+
46+
47+ def test_update_metadata_without_existing_output_namespace ():
48+ """Test that update_metadata sets outputNamespace when not already set."""
49+ # Setup
50+ team_name = "test_team"
51+ team_dict = {
52+ "default" : Team (outputNamespace = "default_namespace" ),
53+ team_name : Team (outputNamespace = "team_namespace" ),
54+ }
55+
56+ # Test without existing outputNamespace
57+ obj = GroupBy (metaData = MetaData (
58+ team = team_name ,
59+ name = "test.group_by.name" ,
60+ ))
61+
62+ # Call the function
63+ parse_teams .update_metadata (obj , team_dict )
64+
65+ # Verify outputNamespace was set from team
66+ assert obj .metaData .outputNamespace == "team_namespace"
67+
68+
69+ def test_update_metadata_preserves_join_part_namespace ():
70+ """Test that update_metadata preserves outputNamespace in join parts."""
71+ # Setup
72+ team_name = "test_team"
73+ team_dict = {
74+ "default" : Team (outputNamespace = "default_namespace" ),
75+ team_name : Team (outputNamespace = "team_namespace" ),
76+ }
77+
78+ # Create a join with join parts that have existing outputNamespace
79+ join_part_gb = GroupBy (metaData = MetaData (outputNamespace = "existing_jp_namespace" ))
80+ join_part = JoinPart (groupBy = join_part_gb )
81+
82+ # Create a join with label parts that have existing outputNamespace
83+ label_part_gb = GroupBy (metaData = MetaData (outputNamespace = "existing_label_namespace" ))
84+ label_parts = LabelParts (labels = [label_part_gb ])
85+
86+ # Create the join object
87+ join = Join (
88+ metaData = MetaData (
89+ team = team_name ,
90+ name = "test.join.name" ,
91+ outputNamespace = "join_namespace"
92+ ),
93+ joinParts = [join_part ],
94+ labelParts = label_parts
95+ )
96+
97+ # Call the function
98+ parse_teams .update_metadata (join , team_dict )
99+
100+ # Verify outputNamespace values were preserved
101+ assert join .metaData .outputNamespace == "join_namespace"
102+ assert join .joinParts [0 ].groupBy .metaData .outputNamespace == "existing_jp_namespace"
103+ assert join .labelParts .labels [0 ].metaData .outputNamespace == "existing_label_namespace"
104+
105+
106+ def test_update_metadata_sets_missing_join_part_namespace ():
107+ """Test that update_metadata sets outputNamespace for join parts when not set."""
108+ # Setup
109+ team_name = "test_team"
110+ team_dict = {
111+ "default" : Team (outputNamespace = "default_namespace" ),
112+ team_name : Team (outputNamespace = "team_namespace" ),
113+ }
114+
115+ # Create a join with join parts that don't have outputNamespace
116+ join_part_gb = GroupBy (metaData = MetaData ())
117+ join_part = JoinPart (groupBy = join_part_gb )
118+
119+ # Create a join with label parts that don't have outputNamespace
120+ label_part_gb = GroupBy (metaData = MetaData ())
121+ label_parts = LabelParts (labels = [label_part_gb ])
122+
123+ # Create the join object
124+ join = Join (
125+ metaData = MetaData (
126+ team = team_name ,
127+ name = "test.join.name" ,
128+ outputNamespace = "join_namespace"
129+ ),
130+ joinParts = [join_part ],
131+ labelParts = label_parts
132+ )
133+
134+ # Call the function
135+ parse_teams .update_metadata (join , team_dict )
136+
137+ # Verify outputNamespace values were set correctly
138+ assert join .metaData .outputNamespace == "join_namespace"
139+ assert join .joinParts [0 ].groupBy .metaData .outputNamespace == "join_namespace"
140+ assert join .labelParts .labels [0 ].metaData .outputNamespace == "join_namespace"
0 commit comments