@@ -618,8 +618,11 @@ def test_check_if_workspace_needs_to_be_terminated_returns_dry_run_is_dry_run_tr
618
618
"terminateUnusedWorkspaces" : "Dry Run" ,
619
619
}
620
620
workspace_id = "123qwe123qwe"
621
+ billable_time = 0
621
622
workspace_helper = workspaces_helper .WorkspacesHelper (session , settings )
622
- result = workspace_helper .check_if_workspace_needs_to_be_terminated (workspace_id )
623
+ result = workspace_helper .check_if_workspace_needs_to_be_terminated (
624
+ workspace_id , billable_time
625
+ )
623
626
assert result == "Yes - Dry Run"
624
627
625
628
@@ -639,8 +642,11 @@ def test_check_if_workspace_needs_to_be_terminated_returns_dry_run_is_dry_run_fa
639
642
"terminateUnusedWorkspaces" : "Dry Run" ,
640
643
}
641
644
workspace_id = "123qwe123qwe"
645
+ billable_time = 0
642
646
workspace_helper = workspaces_helper .WorkspacesHelper (session , settings )
643
- result = workspace_helper .check_if_workspace_needs_to_be_terminated (workspace_id )
647
+ result = workspace_helper .check_if_workspace_needs_to_be_terminated (
648
+ workspace_id , billable_time
649
+ )
644
650
assert result == "Yes - Dry Run"
645
651
646
652
@@ -660,8 +666,11 @@ def test_check_if_workspace_needs_to_be_terminated_returns_empty_string_is_dry_r
660
666
"terminateUnusedWorkspaces" : "Yes" ,
661
667
}
662
668
workspace_id = "123qwe123qwe"
669
+ billable_time = 0
663
670
workspace_helper = workspaces_helper .WorkspacesHelper (session , settings )
664
- result = workspace_helper .check_if_workspace_needs_to_be_terminated (workspace_id )
671
+ result = workspace_helper .check_if_workspace_needs_to_be_terminated (
672
+ workspace_id , billable_time
673
+ )
665
674
assert result == ""
666
675
667
676
@@ -694,10 +703,13 @@ def test_check_if_workspace_needs_to_be_terminated_returns_yes_is_dry_run_false(
694
703
},
695
704
}
696
705
workspace_id = "123qwe123qwe"
706
+ billable_time = 0
697
707
workspace_helper = workspaces_helper .WorkspacesHelper (session , settings )
698
708
mocker .patch .object (workspace_helper , "terminate_unused_workspace" )
699
709
workspace_helper .terminate_unused_workspace .return_value = "Yes"
700
- result = workspace_helper .check_if_workspace_needs_to_be_terminated (workspace_id )
710
+ result = workspace_helper .check_if_workspace_needs_to_be_terminated (
711
+ workspace_id , billable_time
712
+ )
701
713
assert result == "Yes"
702
714
703
715
@@ -735,6 +747,31 @@ def test_check_if_workspace_used_for_selected_period_returns_true_if_timestamp_i
735
747
assert result is True
736
748
737
749
750
+ @freeze_time ("2025-01-15 12:00:00" )
751
+ def test_check_if_workspace_needs_to_be_terminated_returns_empty_string_with_billable_time (
752
+ session ,
753
+ ):
754
+ settings = {
755
+ "region" : "us-east-1" ,
756
+ "hourlyLimits" : 10 ,
757
+ "testEndOfMonth" : "yes" ,
758
+ "isDryRun" : True ,
759
+ "startTime" : 1 ,
760
+ "endTime" : 2 ,
761
+ "terminateUnusedWorkspaces" : "Dry Run" ,
762
+ "dateTimeValues" : {
763
+ "current_month_last_day" : True ,
764
+ },
765
+ }
766
+ workspace_id = "123qwe123qwe"
767
+ billable_time = 5
768
+ workspace_helper = workspaces_helper .WorkspacesHelper (session , settings )
769
+ result = workspace_helper .check_if_workspace_needs_to_be_terminated (
770
+ workspace_id , billable_time
771
+ )
772
+ assert result == ""
773
+
774
+
738
775
def test_get_last_known_user_connection_timestamp_returns_last_connected_time_value (
739
776
session ,
740
777
):
@@ -769,6 +806,50 @@ def test_get_last_known_user_connection_timestamp_returns_last_connected_time_va
769
806
client_stubber .deactivate ()
770
807
771
808
809
+ def test_get_last_known_user_connection_timestamp_handles_pagination (session ):
810
+ settings = {
811
+ "region" : "us-east-1" ,
812
+ "hourlyLimits" : 10 ,
813
+ "testEndOfMonth" : "yes" ,
814
+ "isDryRun" : True ,
815
+ "startTime" : 1 ,
816
+ "endTime" : 2 ,
817
+ }
818
+ workspace_id = "ws-12345678"
819
+ last_known_timestamp = datetime .datetime (
820
+ 2023 , 1 , 1 , 12 , 0 , 0 , tzinfo = datetime .timezone .utc
821
+ )
822
+
823
+ workspace_helper = workspaces_helper .WorkspacesHelper (session , settings )
824
+ client_stubber = Stubber (workspace_helper .workspaces_client )
825
+
826
+ response1 = {"WorkspacesConnectionStatus" : [], "NextToken" : "next-token" }
827
+ expected_params1 = {"WorkspaceIds" : [workspace_id ]}
828
+
829
+ response2 = {
830
+ "WorkspacesConnectionStatus" : [
831
+ {
832
+ "WorkspaceId" : workspace_id ,
833
+ "LastKnownUserConnectionTimestamp" : last_known_timestamp ,
834
+ }
835
+ ]
836
+ }
837
+ expected_params2 = {"WorkspaceIds" : [workspace_id ], "NextToken" : "next-token" }
838
+
839
+ client_stubber .add_response (
840
+ "describe_workspaces_connection_status" , response1 , expected_params1
841
+ )
842
+ client_stubber .add_response (
843
+ "describe_workspaces_connection_status" , response2 , expected_params2
844
+ )
845
+ client_stubber .activate ()
846
+
847
+ result = workspace_helper .get_last_known_user_connection_timestamp (workspace_id )
848
+
849
+ assert result == last_known_timestamp
850
+ client_stubber .deactivate ()
851
+
852
+
772
853
def test_get_last_known_user_connection_timestamp_returns_resource_unavailable_for_exception (
773
854
session ,
774
855
):
@@ -1040,6 +1121,60 @@ def test_get_termination_status_returns_empty_string_when_workspace_not_availabl
1040
1121
assert result == ("" , time .strftime ("%Y-%m-%d" , test_time .timetuple ()))
1041
1122
1042
1123
1124
+ def test_get_termination_status_with_billable_hours (mocker , session ):
1125
+ settings = {
1126
+ "region" : "us-east-1" ,
1127
+ "hourlyLimits" : 10 ,
1128
+ "testEndOfMonth" : True ,
1129
+ "isDryRun" : True ,
1130
+ "startTime" : 1 ,
1131
+ "endTime" : 2 ,
1132
+ "TerminateUnusedWorkspaces" : "Dry Run" ,
1133
+ "dateTimeValues" : {
1134
+ "start_time_for_current_month" : "" ,
1135
+ "end_time_for_current_month" : "" ,
1136
+ "last_day_current_month" : "" ,
1137
+ "first_day_selected_month" : "2025-01-01" ,
1138
+ "start_time_selected_date" : "2025-01-01T00:00:00Z" ,
1139
+ "end_time_selected_date" : "2025-01-02T00:00:00Z" ,
1140
+ "current_month_last_day" : True ,
1141
+ "date_today" : "" ,
1142
+ "date_for_s3_key" : "" ,
1143
+ },
1144
+ }
1145
+ workspace_helper = workspaces_helper .WorkspacesHelper (session , settings )
1146
+
1147
+ workspace_id = "ws-12345678"
1148
+ billable_time = 10
1149
+ tags = [{"Key" : "TestTag" , "Value" : "TestValue" }]
1150
+
1151
+ mock_timestamp = datetime .datetime (2025 , 1 , 15 , 12 , 0 , 0 )
1152
+ mocker .patch .object (
1153
+ workspace_helper ,
1154
+ "get_last_known_user_connection_timestamp" ,
1155
+ return_value = mock_timestamp ,
1156
+ )
1157
+ mocker .patch .object (
1158
+ workspace_utils , "is_terminate_workspace_enabled" , return_value = True
1159
+ )
1160
+ mocker .patch .object (
1161
+ workspace_helper , "check_if_workspace_available_on_first_day_selected_month"
1162
+ )
1163
+ mocker .patch .object (workspace_utils , "check_if_workspace_used_for_selected_period" )
1164
+ mocker .patch .object (workspace_helper , "check_if_workspace_needs_to_be_terminated" )
1165
+
1166
+ result = workspace_helper .get_termination_status (workspace_id , billable_time , tags )
1167
+
1168
+ assert result == ("" , "2025-01-15" )
1169
+ workspace_helper .get_last_known_user_connection_timestamp .assert_called_once_with (
1170
+ workspace_id
1171
+ )
1172
+ workspace_utils .is_terminate_workspace_enabled .assert_called_once ()
1173
+ workspace_helper .check_if_workspace_available_on_first_day_selected_month .assert_called_once ()
1174
+ workspace_utils .check_if_workspace_used_for_selected_period .assert_called_once ()
1175
+ workspace_helper .check_if_workspace_needs_to_be_terminated .assert_not_called ()
1176
+
1177
+
1043
1178
def test_get_workspaces_for_directory_use_next_token (session ):
1044
1179
settings = {
1045
1180
"region" : "us-east-1" ,
@@ -1486,8 +1621,11 @@ def test_check_if_workspace_needs_to_be_terminated_returns_empty_string_for_last
1486
1621
},
1487
1622
}
1488
1623
workspace_id = "123qwe123qwe"
1624
+ billable_time = 0
1489
1625
workspace_helper = workspaces_helper .WorkspacesHelper (session , settings )
1490
- result = workspace_helper .check_if_workspace_needs_to_be_terminated (workspace_id )
1626
+ result = workspace_helper .check_if_workspace_needs_to_be_terminated (
1627
+ workspace_id , billable_time
1628
+ )
1491
1629
assert result == ""
1492
1630
1493
1631
0 commit comments