Skip to content

Commit 7fd742a

Browse files
vasilylvasilyl
committed
Add new table initialization and error handling unit tests for BigQuery analytics plugin
1 parent b9ee0a3 commit 7fd742a

1 file changed

Lines changed: 272 additions & 0 deletions

File tree

‎plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go‎

Lines changed: 272 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -388,3 +388,275 @@ func TestSchemaFieldsMatch(t *testing.T) {
388388
t.Errorf("Expected second sub-field of RECORD to be 'sub2', got %s", updatedRecord.Schema[1].Name)
389389
}
390390
}
391+
392+
func TestNewBigQueryAgentAnalyticsPlugin_TableExists(t *testing.T) {
393+
ctx := context.Background()
394+
config := DefaultConfig()
395+
config.Enabled = true
396+
config.ProjectID = "test-project"
397+
config.DatasetID = "test-dataset"
398+
config.TableName = "test-table"
399+
config.CreateViews = false
400+
config.AutoSchemaUpgrade = false
401+
402+
mockTransport := &mockTransport{
403+
roundTrip: func(r *http.Request) (*http.Response, error) {
404+
if r.Method == "GET" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables/test-table") {
405+
return &http.Response{
406+
StatusCode: http.StatusOK,
407+
Body: io.NopCloser(strings.NewReader(`{"tableReference":{"projectId":"test-project","datasetId":"test-dataset","tableId":"test-table"}}`)),
408+
}, nil
409+
}
410+
return &http.Response{
411+
StatusCode: http.StatusOK,
412+
Body: io.NopCloser(strings.NewReader("{}")),
413+
}, nil
414+
},
415+
}
416+
httpClient := &http.Client{Transport: mockTransport}
417+
bqClient, err := bq.NewClient(ctx, config.ProjectID, option.WithHTTPClient(httpClient))
418+
if err != nil {
419+
t.Fatalf("Failed to create bigquery client: %v", err)
420+
}
421+
422+
lis, err := net.Listen("tcp", "localhost:0")
423+
if err != nil {
424+
t.Fatalf("failed to listen: %v", err)
425+
}
426+
gSrv := grpc.NewServer()
427+
storagepb.RegisterBigQueryWriteServer(gSrv, &fakeBigQueryWriteServer{})
428+
go gSrv.Serve(lis)
429+
t.Cleanup(gSrv.Stop)
430+
431+
conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
432+
if err != nil {
433+
t.Fatalf("failed to dial test server: %v", err)
434+
}
435+
t.Cleanup(func() { conn.Close() })
436+
437+
writeClient, err := bqstorage.NewBigQueryWriteClient(ctx, option.WithGRPCConn(conn))
438+
if err != nil {
439+
t.Fatalf("Failed to create BigQuery write client: %v", err)
440+
}
441+
442+
p, err := NewBigQueryAgentAnalyticsPluginWithClients(ctx, config, bqClient, writeClient)
443+
if err != nil {
444+
t.Fatalf("Expected no error when table exists, got: %v", err)
445+
}
446+
if p == nil {
447+
t.Fatal("Expected plugin to be non-nil")
448+
}
449+
}
450+
451+
func TestNewBigQueryAgentAnalyticsPlugin_TableNotFound_CreateSucceeds(t *testing.T) {
452+
ctx := context.Background()
453+
config := DefaultConfig()
454+
config.Enabled = true
455+
config.ProjectID = "test-project"
456+
config.DatasetID = "test-dataset"
457+
config.TableName = "test-table"
458+
config.CreateViews = false
459+
config.AutoSchemaUpgrade = false
460+
461+
createCalled := false
462+
463+
mockTransport := &mockTransport{
464+
roundTrip: func(r *http.Request) (*http.Response, error) {
465+
if r.Method == "GET" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables/test-table") {
466+
return &http.Response{
467+
StatusCode: http.StatusNotFound,
468+
Body: io.NopCloser(strings.NewReader(`{"error":{"code":404,"message":"Not found"}}`)),
469+
}, nil
470+
}
471+
if r.Method == "POST" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables") {
472+
createCalled = true
473+
return &http.Response{
474+
StatusCode: http.StatusOK,
475+
Body: io.NopCloser(strings.NewReader("{}")),
476+
}, nil
477+
}
478+
return &http.Response{
479+
StatusCode: http.StatusOK,
480+
Body: io.NopCloser(strings.NewReader("{}")),
481+
}, nil
482+
},
483+
}
484+
httpClient := &http.Client{Transport: mockTransport}
485+
bqClient, err := bq.NewClient(ctx, config.ProjectID, option.WithHTTPClient(httpClient))
486+
if err != nil {
487+
t.Fatalf("Failed to create bigquery client: %v", err)
488+
}
489+
490+
lis, err := net.Listen("tcp", "localhost:0")
491+
if err != nil {
492+
t.Fatalf("failed to listen: %v", err)
493+
}
494+
gSrv := grpc.NewServer()
495+
storagepb.RegisterBigQueryWriteServer(gSrv, &fakeBigQueryWriteServer{})
496+
go gSrv.Serve(lis)
497+
t.Cleanup(gSrv.Stop)
498+
499+
conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
500+
if err != nil {
501+
t.Fatalf("failed to dial test server: %v", err)
502+
}
503+
t.Cleanup(func() { conn.Close() })
504+
505+
writeClient, err := bqstorage.NewBigQueryWriteClient(ctx, option.WithGRPCConn(conn))
506+
if err != nil {
507+
t.Fatalf("Failed to create BigQuery write client: %v", err)
508+
}
509+
510+
p, err := NewBigQueryAgentAnalyticsPluginWithClients(ctx, config, bqClient, writeClient)
511+
if err != nil {
512+
t.Fatalf("Expected no error when table is created successfully, got: %v", err)
513+
}
514+
if p == nil {
515+
t.Fatal("Expected plugin to be non-nil")
516+
}
517+
if !createCalled {
518+
t.Error("Expected table creation to be attempted, but it was not called")
519+
}
520+
}
521+
522+
func TestNewBigQueryAgentAnalyticsPlugin_TableNotFound_CreateFails(t *testing.T) {
523+
ctx := context.Background()
524+
config := DefaultConfig()
525+
config.Enabled = true
526+
config.ProjectID = "test-project"
527+
config.DatasetID = "test-dataset"
528+
config.TableName = "test-table"
529+
config.CreateViews = false
530+
config.AutoSchemaUpgrade = false
531+
532+
mockTransport := &mockTransport{
533+
roundTrip: func(r *http.Request) (*http.Response, error) {
534+
if r.Method == "GET" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables/test-table") {
535+
return &http.Response{
536+
StatusCode: http.StatusNotFound,
537+
Body: io.NopCloser(strings.NewReader(`{"error":{"code":404,"message":"Not found"}}`)),
538+
}, nil
539+
}
540+
if r.Method == "POST" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables") {
541+
return &http.Response{
542+
StatusCode: http.StatusForbidden,
543+
Body: io.NopCloser(strings.NewReader(`{"error":{"code":403,"message":"Permission denied"}}`)),
544+
}, nil
545+
}
546+
return &http.Response{
547+
StatusCode: http.StatusOK,
548+
Body: io.NopCloser(strings.NewReader("{}")),
549+
}, nil
550+
},
551+
}
552+
httpClient := &http.Client{Transport: mockTransport}
553+
bqClient, err := bq.NewClient(ctx, config.ProjectID, option.WithHTTPClient(httpClient))
554+
if err != nil {
555+
t.Fatalf("Failed to create bigquery client: %v", err)
556+
}
557+
558+
lis, err := net.Listen("tcp", "localhost:0")
559+
if err != nil {
560+
t.Fatalf("failed to listen: %v", err)
561+
}
562+
gSrv := grpc.NewServer()
563+
storagepb.RegisterBigQueryWriteServer(gSrv, &fakeBigQueryWriteServer{})
564+
go gSrv.Serve(lis)
565+
t.Cleanup(gSrv.Stop)
566+
567+
conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
568+
if err != nil {
569+
t.Fatalf("failed to dial test server: %v", err)
570+
}
571+
t.Cleanup(func() { conn.Close() })
572+
573+
writeClient, err := bqstorage.NewBigQueryWriteClient(ctx, option.WithGRPCConn(conn))
574+
if err != nil {
575+
t.Fatalf("Failed to create BigQuery write client: %v", err)
576+
}
577+
578+
p, err := NewBigQueryAgentAnalyticsPluginWithClients(ctx, config, bqClient, writeClient)
579+
if err == nil {
580+
t.Fatal("Expected error when table creation fails, got nil")
581+
}
582+
if p != nil {
583+
t.Error("Expected plugin to be nil when table creation fails")
584+
}
585+
if !strings.Contains(err.Error(), "failed to create BigQuery table") {
586+
t.Errorf("Expected error message to contain 'failed to create BigQuery table', got: %v", err)
587+
}
588+
}
589+
590+
func TestNewBigQueryAgentAnalyticsPlugin_MetadataError_Non404(t *testing.T) {
591+
ctx := context.Background()
592+
config := DefaultConfig()
593+
config.Enabled = true
594+
config.ProjectID = "test-project"
595+
config.DatasetID = "test-dataset"
596+
config.TableName = "test-table"
597+
config.CreateViews = false
598+
config.AutoSchemaUpgrade = false
599+
600+
createAttempted := false
601+
602+
mockTransport := &mockTransport{
603+
roundTrip: func(r *http.Request) (*http.Response, error) {
604+
if r.Method == "GET" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables/test-table") {
605+
return &http.Response{
606+
StatusCode: http.StatusForbidden,
607+
Body: io.NopCloser(strings.NewReader(`{"error":{"code":403,"message":"Permission denied"}}`)),
608+
}, nil
609+
}
610+
if r.Method == "POST" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables") {
611+
createAttempted = true
612+
return &http.Response{
613+
StatusCode: http.StatusOK,
614+
Body: io.NopCloser(strings.NewReader("{}")),
615+
}, nil
616+
}
617+
return &http.Response{
618+
StatusCode: http.StatusOK,
619+
Body: io.NopCloser(strings.NewReader("{}")),
620+
}, nil
621+
},
622+
}
623+
httpClient := &http.Client{Transport: mockTransport}
624+
bqClient, err := bq.NewClient(ctx, config.ProjectID, option.WithHTTPClient(httpClient))
625+
if err != nil {
626+
t.Fatalf("Failed to create bigquery client: %v", err)
627+
}
628+
629+
lis, err := net.Listen("tcp", "localhost:0")
630+
if err != nil {
631+
t.Fatalf("failed to listen: %v", err)
632+
}
633+
gSrv := grpc.NewServer()
634+
storagepb.RegisterBigQueryWriteServer(gSrv, &fakeBigQueryWriteServer{})
635+
go gSrv.Serve(lis)
636+
t.Cleanup(gSrv.Stop)
637+
638+
conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
639+
if err != nil {
640+
t.Fatalf("failed to dial test server: %v", err)
641+
}
642+
t.Cleanup(func() { conn.Close() })
643+
644+
writeClient, err := bqstorage.NewBigQueryWriteClient(ctx, option.WithGRPCConn(conn))
645+
if err != nil {
646+
t.Fatalf("Failed to create BigQuery write client: %v", err)
647+
}
648+
649+
p, err := NewBigQueryAgentAnalyticsPluginWithClients(ctx, config, bqClient, writeClient)
650+
if err == nil {
651+
t.Fatal("Expected error when metadata retrieval fails with non-404 error, got nil")
652+
}
653+
if p != nil {
654+
t.Error("Expected plugin to be nil when metadata retrieval fails")
655+
}
656+
if createAttempted {
657+
t.Error("Expected table creation NOT to be attempted, but it was called")
658+
}
659+
if !strings.Contains(err.Error(), "failed to retrieve BigQuery table metadata") {
660+
t.Errorf("Expected error message to contain 'failed to retrieve BigQuery table metadata', got: %v", err)
661+
}
662+
}

0 commit comments

Comments
 (0)