Skip to content

Commit

Permalink
test(codegen): add codegen util testing
Browse files Browse the repository at this point in the history
jira: CORP-5480
  • Loading branch information
Mandeep Bahal committed Sep 10, 2024
1 parent 8e77443 commit c9da754
Show file tree
Hide file tree
Showing 6 changed files with 152 additions and 46 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@

![Actions Status](https://github.com/DeloitteOptimalReality/airflow-wrapper-go/actions/workflows/Linux.yml/badge.svg?branch=CORP-5480-update-dag-template)
[![Actions Status](https://github.com/DeloitteOptimalReality/airflow-wrapper-go/actions/workflows/Linux.yml/badge.svg?branch=CORP-5480-update-dag-template)](https://github.com/DeloitteOptimalReality/airflow-wrapper-go/actions)

# Introduction
This is a wrapper library of the airflow-client-go library. The library is written in Go and allows your to create clients with interface methods to interact with the Apache Airflow REST API.
Expand Down
4 changes: 4 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,12 @@ require (
)

require (
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/golang/protobuf v1.4.2 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/stretchr/testify v1.9.0 // indirect
golang.org/x/oauth2 v0.0.0-20210218202405-ba52d332ba99 // indirect
google.golang.org/appengine v1.6.6 // indirect
google.golang.org/protobuf v1.25.0 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
7 changes: 7 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,8 @@ github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMn
github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98=
Expand Down Expand Up @@ -105,11 +107,14 @@ github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+o
github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
Expand Down Expand Up @@ -352,6 +357,8 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
honnef.co/go/tools v0.0.0-20190106161140-3f1c8253044a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
Expand Down
37 changes: 19 additions & 18 deletions pkg/codegen/codegen.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,12 @@ type Args struct {
}

// Function to convert an interface to a JSON string
func toJSONString(v interface{}) string {
func ToJSONString(v interface{}) (string, error) {
jsonData, err := json.Marshal(v)
if err != nil {
return ""
return "", err
}
return string(jsonData)
return string(jsonData), nil
}

type StartDate struct {
Expand Down Expand Up @@ -66,7 +66,7 @@ type GenData struct {
Tasks []HttpOperator
}

func checkDeps(deps []string) bool {
func CheckDeps(deps []string) bool {
return len(deps) > 0
}

Expand All @@ -79,17 +79,18 @@ func writeValueToBuffer(v interface{}, buf *bytes.Buffer) bool {
buf.WriteString("False")
}
case string:
buf.WriteString("'")
buf.WriteString("\"")
buf.WriteString(v)
buf.WriteString("'")
buf.WriteString("\"")
default:
buf.WriteString(fmt.Sprintf("%v", v))
}
return true
}

func toMap(v interface{}) (map[string]interface{}, error) {
func ToMap(v interface{}) (map[string]interface{}, error) {
var res map[string]interface{}
fmt.Println("MARSHAlling")
a, err := json.Marshal(v)
if err != nil {
return res, err
Expand All @@ -99,7 +100,7 @@ func toMap(v interface{}) (map[string]interface{}, error) {
}

// Function to convert a Go map to a Python dictionary string
func mapToPythonDict(m map[string]interface{}) (string, error) {
func MapToPythonDict(m map[string]interface{}) (string, error) {
var buf bytes.Buffer
buf.WriteString("{")
first := true
Expand All @@ -108,9 +109,9 @@ func mapToPythonDict(m map[string]interface{}) (string, error) {
buf.WriteString(", ")
}
first = false
buf.WriteString("'")
buf.WriteString("\"")
buf.WriteString(k)
buf.WriteString("': ")
buf.WriteString("\": ")
switch v := v.(type) {
case []interface{}:
buf.WriteString("[")
Expand All @@ -119,7 +120,7 @@ func mapToPythonDict(m map[string]interface{}) (string, error) {
}
buf.WriteString("]")
case map[string]interface{}:
nested, err := mapToPythonDict(v)
nested, err := MapToPythonDict(v)
if err != nil {
return "", err
}
Expand All @@ -140,7 +141,7 @@ func BoolTitle(b bool) string {
}

// Function to transform the task ID
func transformTaskID(taskID string) string {
func TransformTaskID(taskID string) string {
return "wt_" + strings.ReplaceAll(taskID, "-", "_")
}

Expand All @@ -153,17 +154,17 @@ func CreateDagGen(g GenData, directory string) (string, error) {
// Prepare a map of original task IDs to transformed task IDs
taskIDMap := make(map[string]string)
for _, task := range data.Tasks {
taskIDMap[task.TaskID] = transformTaskID(task.TaskID)
taskIDMap[task.TaskID] = TransformTaskID(task.TaskID)
}

t := template.New("dag").Funcs(
template.FuncMap{
"toJSONString": toJSONString,
"toMap": toMap,
"toJSONString": ToJSONString,
"toMap": ToMap,
"BoolTitle": BoolTitle,
"mapToPythonDict": mapToPythonDict,
"checkDeps": checkDeps,
"transformTaskID": transformTaskID,
"mapToPythonDict": MapToPythonDict,
"checkDeps": CheckDeps,
"transformTaskID": TransformTaskID,
"originalTaskIDMap": func() map[string]string { return taskIDMap },
})
tp, err := t.Parse(tmpl)
Expand Down
27 changes: 0 additions & 27 deletions pkg/codegen/codegen_test.go

This file was deleted.

121 changes: 121 additions & 0 deletions test/codegen/codegen_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,121 @@
package codegen

import (
"fmt"
"testing"

"github.com/DeloitteOptimalReality/airflow-wrapper-go/pkg/codegen"
"github.com/stretchr/testify/assert"
)

func TestCreateDagObject(t *testing.T) {
tests := []struct {
name string
want string
}{
{"test1", ""},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
data := codegen.GenData{DagDef: codegen.Dag{ID: ""}}
if got, err := codegen.CreateDagGen(data, ""); false {
if err != nil {
t.Error()
}
fmt.Println(got)
t.Errorf("CreateDagObject() = %v, want %v", got, tt.want)
}
})
}
}

func TestToMapValid(t *testing.T) {
testStruct := codegen.HttpOperator{
TaskID: "test_task_id",
ConnectionId: "test_connection_id",
Name: "test_name",
Endpoint: "test_endpoint",
Data: "test_data",
Downstream: []string{"downstream_task_1", "downstream_task_2"},
}
res, err := codegen.ToMap(testStruct)
if err != nil {
t.Error(err)
}
assert.Equal(t, res["Name"], "test_name")
assert.Equal(t, res["ConnectionId"], "test_connection_id")
assert.Equal(t, res["Name"], "test_name")
assert.Equal(t, res["Endpoint"], "test_endpoint")
assert.Equal(t, res["Data"], "test_data")
assert.ElementsMatch(t, res["Downstream"], []string{"downstream_task_1", "downstream_task_2"})
}

// TODO: Test the invalid case for ToMap function. I'm not sure how to hit the edge case yet.

func TestMapToPythonDict(t *testing.T) {
expectedPythonDict := `
{
"key1": {
"ID": "dag_id",
"Description": "dag_description",
"StartDate": {
"Day": 1,
"Month": 1,
"Year": 2024
},
"Tags": [],
"DefaultArgs": {
"Email": "test_email",
"Retries": 1,
"RetryDelay": 999
}
}
}
`

testMap := map[string]interface{}{
"key1": map[string]interface{}{
"ID": "dag_id",
"Description": "dag_description",
"StartDate": map[string]interface{}{
"Day": 1,
"Month": 1,
"Year": 2024,
},
"Tags": []string{},
"DefaultArgs": map[string]interface{}{
"Email": "test_email",
"Retries": 1,
"RetryDelay": 999,
},
},
}
res, err := codegen.MapToPythonDict(testMap)
if err != nil {
t.Error(err)
}
assert.JSONEq(t, res, expectedPythonDict)
}

func TestToJSONString(t *testing.T) {
expectedJSONString := `
{
"key1": "value1"
}
`

testMap := map[string]interface{}{
"key1": "value1",
}
res, err := codegen.ToJSONString(testMap)
if err != nil {
t.Error(err)
}
assert.JSONEq(t, res, expectedJSONString)
}

func TestCheckDeps(t *testing.T) {
assert.True(t, codegen.CheckDeps([]string{"dep_1"}))
assert.True(t, codegen.CheckDeps([]string{"dep_1", "dep_2"}))
assert.False(t, codegen.CheckDeps([]string{}))
}

0 comments on commit c9da754

Please sign in to comment.